|
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM |
|
import gradio as gr |
|
|
|
model_name = "google/flan-t5-base" |
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
model = AutoModelForSeq2SeqLM.from_pretrained(model_name) |
|
|
|
def text_2_text_generation(input): |
|
|
|
input_ids = tokenizer(input, return_tensors="pt", truncation= True).input_ids |
|
|
|
output = model.generate(input_ids, |
|
max_new_tokens = 200, |
|
do_sample = True, |
|
top_p = 0.9, |
|
top_k = 50) |
|
|
|
output_text = tokenizer.decode(output[0], skip_special_tokens=True) |
|
|
|
|
|
return output_text |
|
|
|
iface = gr.Interface( |
|
fn = text_2_text_generation, |
|
inputs = gr.Textbox(label = "Enter your queries...", lines = 5), |
|
outputs = gr.Textbox(lines = 5) |
|
) |
|
|
|
iface.launch() |