File size: 508 Bytes
3b25749 2ff483e 3b25749 2ff483e 3b25749 2ff483e |
1 2 3 4 5 6 7 8 9 10 11 12 13 |
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() |