Spaces:
Sleeping
Sleeping
import gradio | |
import torch | |
from transformers import AutoModelForCausalLM, AutoTokenizer | |
model = AutoModelForCausalLM.from_pretrained("cyberagent/open-calm-small", device_map="auto", torch_dtype=torch.float16) | |
tokenizer = AutoTokenizer.from_pretrained("cyberagent/open-calm-small") | |
def my_inference_function(name): | |
mail="お世話になっております。商品開発部の森本です。先日は貴重な" | |
inputs = tokenizer(f"(上司へのメール敬語で){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(f"(上司へのメール敬語で){mail}","") | |
output=output.replace("\n", "") | |
return output | |
gradio_interface = gradio.Interface( | |
fn = my_inference_function, | |
inputs = "text", | |
outputs = "text" | |
) | |
gradio_interface.launch() |