| import torch | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline | |
| import gradio as gr | |
| model_name = "Writer/palmyra-small" | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model = AutoModelForCausalLM.from_pretrained(model_name).to(device) | |
| text_generator = TextGenerationPipeline(model=model, tokenizer=tokenizer) | |
| iface = gr.Interface(fn=text_generator, inputs="text", outputs="text") | |
| iface.launch() |