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()