ttt / app.py
akoksal's picture
update
10a2445
raw
history blame
1.27 kB
import gradio as gr
from transformers import AutoTokenizer, pipeline
import torch
tokenizer = AutoTokenizer.from_pretrained("notexist/ttt")
tdk = pipeline('text-generation', model='notexist/ttt', tokenizer=tokenizer)
def predict(name, sl, topk, topp):
if name == "":
x = tdk(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 x[len(f"<|endoftext|>"):]
else:
x = tdk(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 x[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()