import os
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

hf_token= os.getenv("access_token")
tokenizer = AutoTokenizer.from_pretrained("afrizalha/Sasando-1-25M", token=hf_token)
tiny = AutoModelForCausalLM.from_pretrained("afrizalha/Sasando-1-25M", token=hf_token)
tinier = AutoModelForCausalLM.from_pretrained("afrizalha/Sasando-1-7M", token=hf_token)


desc = """Sasando-1 is a tiny, highly experimental text generator built using the Phi-3 architecture. It comes with two variations of microscopic sizes: 7M and 25M parameters. It is trained on a tightly-controlled Indo4B dataset filtered to only have 18000 unique words. The method is inspired by Microsoft's TinyStories paper which demonstrates that a tiny language model can produce fluent text when trained on tightly-controlled dataset.\n\nTry prompting with two simple words, and let the model continue. Fun examples provided below."""

def generate(starting_text, choice, temp, top_p):
    if choice == '7M':
        model = tinier
    elif choice == '25M':
        model = tiny
    elif choice == "Info":
        yield desc
        return
    
    results = []
    for i in range(5):
        inputs = tokenizer([starting_text], return_tensors="pt").to(model.device)
        outputs = model.generate(
            inputs=inputs.input_ids,
            max_new_tokens=32-len(inputs.input_ids[0]),
            do_sample=True,
            temperature=temp,
            top_p=top_p
        )
        outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
        outputs = outputs[:outputs.find(".")]
        results.append(outputs)
        yield "\n\n".join(results)

with gr.Blocks(theme=gr.themes.Soft()) as app:
    starting_text = gr.Textbox(label="Starting text", value="cinta adalah")
    res = gr.Textbox(label="Continuation", value="cinta adalah", scale=2)
    choice = gr.Radio(["7M", "25M", "Info"], label="Select model", value='Info')
    temp = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, step=0.1, value=0.7)
    top_p = gr.Slider(label="Top P", minimum=0.1, maximum=1.0, step=0.1, value=0.5)
    gr.Interface(
        fn=generate,
         inputs=[starting_text,choice,temp,top_p],
         outputs=[res],
         allow_flagging="never",
         title="Sasando-1",
         )
    examples=gr.Examples([["gue"], ["presiden"], ["cinta adalah"], ["allah, aku"], ["dia marah karena"],
                       ["inflasi"], ["kolam renang"], ["messi"], ["jalan-jalan"], ["komputer itu"]], [starting_text])
app.launch()