auto_completion / app.py
KennethTM's picture
added app file
2c38237
import gradio as gr
from threading import Thread
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
model_id = "KennethTM/gpt-neo-1.3B-danish"
peft_model_id = "KennethTM/gpt-neo-1.3B-danish-lora-r64-a128"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
model.load_adapter(peft_model_id)
model.eval()
def run_generation(user_text):
model_inputs = tokenizer(user_text, return_tensors="pt")
streamer = TextIteratorStreamer(tokenizer, skip_prompt=False, skip_special_tokens=True)
generate_kwargs = dict(
inputs = model_inputs.input_ids,
attention_mask = model_inputs.attention_mask,
streamer = streamer,
max_new_tokens = 20,
penalty_alpha = 0.6,
top_k = 4,
)
t = Thread(target=model.generate, kwargs=generate_kwargs)
t.start()
model_output = ""
for new_text in streamer:
model_output += new_text
yield model_output
return model_output
with gr.Blocks() as demo:
gr.Markdown("# Auto-completion med dansk sprogmodel")
gr.Markdown("## Tryk 'Kør' eller 'SHIFT+ENTER' for auto-completion!")
textbox = gr.Textbox(placeholder="Skriv her...", lines=10)
btn = gr.Button("Kør")
btn.click(run_generation, textbox, textbox)
textbox.submit(run_generation, textbox, textbox)
gr.Markdown("## Om")
gr.Markdown("*Dansk sprogmodel (KennethTM/gpt-neo-1.3B-danish) fin-tunet yderligere med PEFT på udvalgte tekster om natur og miljø*")
demo.queue().launch()