t0-generation / app.py
ttj's picture
Use one model first
84b2150
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
def get_pipe(name):
tokenizer = AutoTokenizer.from_pretrained(name)
model = AutoModelForSeq2SeqLM.from_pretrained(name)
pipe = pipeline(
"summarization", model=model, tokenizer=tokenizer, framework="pt"
)
return pipe
model_names = ['bigscience/T0_3B'] #, 'bigscience/T0p', 'bigscience/T0pp']
#model_names = ['bigscience/T0_3B','bigscience/T0'] #, 'bigscience/T0p', 'bigscience/T0pp']
pipes = [get_pipe(name) for name in model_names]
def _fn(text, do_sample, min_length, max_length, temperature, top_p, pipe):
out = pipe(
text,
do_sample=do_sample,
min_length=min_length,
max_length=max_length,
temperature=temperature,
top_p=top_p,
truncation=True,
)
return out[0]["summary_text"]
def fn(*args):
return [_fn(*args, pipe=pipe) for pipe in pipes]
import gradio as gr
interface = gr.Interface(
fn,
inputs=[
gr.inputs.Textbox(lines=10, label="input text"),
gr.inputs.Checkbox(label="do_sample", default=True),
gr.inputs.Slider(1, 128, step=1, default=64, label="min_length"),
gr.inputs.Slider(1, 128, step=1, default=64, label="max_length"),
gr.inputs.Slider(0.0, 1.0, step=0.1, default=1, label="temperature"),
gr.inputs.Slider(0.0, 1.0, step=0.1, default=1, label="top_p"),
],
outputs=[
gr.outputs.Textbox(label=f"output by {name}") for name in model_names
],
#examples=[[ex] for ex in examples],
title="T0 playground",
description="""
This is a playground for playing around with T0 models.
See https://huggingface.co/bigscience/T0 for more details
""",
)
interface.launch()