Spaces:
Sleeping
Sleeping
import gradio | |
import torch | |
from transformers import AutoModelForCausalLM, AutoTokenizer | |
model = AutoModelForCausalLM.from_pretrained("cyberagent/open-calm-medium", device_map="auto", torch_dtype=torch.float32) | |
tokenizer = AutoTokenizer.from_pretrained("cyberagent/open-calm-medium") | |
def my_inference_function(mail): | |
mail=mail.replace("\n", "") | |
inputs = tokenizer(mail, return_tensors="pt").to(model.device) | |
with torch.no_grad(): | |
tokens = model.generate( | |
**inputs, | |
max_new_tokens=4, | |
do_sample=True, | |
temperature=0.7, | |
top_p=0.9, | |
repetition_penalty=1.05, | |
pad_token_id=tokenizer.pad_token_id, | |
) | |
output = tokenizer.decode(tokens[0], skip_special_tokens=True) | |
output=output.replace("(上司へのメール敬語で)","") | |
output=output.replace(mail,"") | |
output=output.replace("\n", "") | |
return output | |
gradio_interface = gradio.Interface( | |
fn = my_inference_function, | |
inputs = "text", | |
outputs = "text" | |
) | |
gradio_interface.launch() |