| import gradio as gr |
|
|
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| |
| |
| tokenizer = AutoTokenizer.from_pretrained("cyberagent/open-calm-3b") |
|
|
| def proc( inputs ): |
| with torch.no_grad(): |
| tokens = model.generate( |
| **inputs, |
| max_new_tokens=64, |
| do_sample=True, |
| temperature=0.7, |
| pad_token_id=tokenizer.pad_token_id, |
| ) |
| |
| return tokenizer.decode(tokens[0], skip_special_tokens=True) |
|
|
| def greet(name): |
| inputs = tokenizer(name, return_tensors="pt").to(model.device) |
| |
| |
| return inputs |
|
|
| iface = gr.Interface(fn=greet, inputs="text", outputs="text") |
| iface.launch() |
|
|