ttt / app.py
akoksal's picture
Update app
7c6863d
raw
history blame
2 kB
import gradio as gr
from transformers import AutoTokenizer, pipeline
import torch
tokenizer1 = AutoTokenizer.from_pretrained("notexist/ttt")
tdk1 = pipeline('text-generation', model='notexist/ttt', tokenizer=tokenizer)
tokenizer2 = AutoTokenizer.from_pretrained("notexist/ttt")
tdk2 = pipeline('text-generation', model='notexist/ttt', tokenizer=tokenizer)
def predict(name, sl, topk, topp):
if name == "":
x1 = tdk1(f"<|endoftext|>",
do_sample=True,
max_length=64,
top_k=topk,
top_p=topp,
num_return_sequences=1,
repetition_penalty=sl
)[0]["generated_text"]
x2 = tdk1(f"<|endoftext|>",
do_sample=True,
max_length=64,
top_k=topk,
top_p=topp,
num_return_sequences=1,
repetition_penalty=sl
)[0]["generated_text"]
return x1[len(f"<|endoftext|>"):]+"\n\n"+x2[len(f"<|endoftext|>"):]
else:
x1 = tdk1(f"<|endoftext|>{name}\n\n",
do_sample=True,
max_length=64,
top_k=topk,
top_p=topp,
num_return_sequences=1,
repetition_penalty=sl
)[0]["generated_text"]
x2 = tdk2(f"<|endoftext|>{name}\n\n",
do_sample=True,
max_length=64,
top_k=topk,
top_p=topp,
num_return_sequences=1,
repetition_penalty=sl
)[0]["generated_text"]
return x1[len(f"<|endoftext|>{name}\n\n"):]+"\n\n"+x2[len(f"<|endoftext|>{name}\n\n"):]
iface = gr.Interface(fn=predict, inputs=["text",\
gr.inputs.Slider(0, 3, default=1.1, label="repetition_penalty"),\
gr.inputs.Slider(0, 100, default=75, label="top_k"),\
gr.inputs.Slider(0, 1, default=0.95, label="top_p")]
, outputs="text")
iface.launch()