| | import gradio as gr |
| | from transformers import AutoModelForSeq2SeqLM, AutoTokenizer |
| | import torch |
| |
|
| | |
| | model_name = "google/flan-t5-small" |
| | tokenizer = AutoTokenizer.from_pretrained(model_name) |
| | model = AutoModelForSeq2SeqLM.from_pretrained(model_name) |
| |
|
| | def chat(prompt): |
| | input_ids = tokenizer(prompt, return_tensors="pt").input_ids |
| | with torch.no_grad(): |
| | outputs = model.generate(input_ids, max_new_tokens=200) |
| | reply = tokenizer.decode(outputs[0], skip_special_tokens=True) |
| | return reply |
| |
|
| | demo = gr.Interface(fn=chat, inputs="text", outputs="text", title="FLAN-T5 Chatbot") |
| |
|
| | demo.launch() |
| |
|