import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from hf_olmo import * # registers the Auto* classes


title = """# 👋🏻Welcome to 🌟Tonic's 👴🏻🏑Olmo-on-Device
"[👴🏻🏑allenai/OLMo-1B](https://huggingface.co/allenai/OLMo-1B) is an on-device LLM from Allen-ai that can fit on your laptop ! You can use this demo to try out their model. You can also use [👴🏻🏑allenai/OLMo-1B](https://huggingface.co/allenai/OLMo-7B) [on your laptop & by cloning this space](https://huggingface.co/spaces/Tonic/Olmo/tree/main?clone=true). 🧬🔬🔍 Simply click here: <a style="display:inline-block" href="https://huggingface.co/spaces/Tonic/Olmo?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></h3> 
Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder's🛠️community 👻  [![Join us on Discord](https://img.shields.io/discord/1109943800132010065?label=Discord&logo=discord&style=flat-square)](https://discord.gg/GWpVpekp) On 🤗Huggingface: [TeamTonic](https://huggingface.co/TeamTonic) & [MultiTransformer](https://huggingface.co/MultiTransformer) On 🌐Github: [Tonic-AI](https://github.com/tonic-ai) & contribute to 🌟 [DataTonic](https://github.com/Tonic-AI/DataTonic) 🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗
"""

model_name = "allenai/OLMo-1B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def generate_text(prompt, max_new_tokens, top_k, top_p, do_sample):
    inputs = tokenizer(prompt, return_tensors='pt', return_token_type_ids=False)
    response = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=do_sample, top_k=top_k, top_p=top_p)
    return tokenizer.batch_decode(response, skip_special_tokens=True)[0]

def main():
    with gr.Blocks() as demo:
        gr.Markdown(title)
        output = gr.Textbox(label="👴🏻🏑OLMo", lines=10)
        submit_button = gr.Button("🌟Try it Out!")
        with gr.Accordion("🤖Optional Parameters", open=False):
            with gr.Row(): 
                do_sample = gr.Checkbox(value=False, label="Untick for faster inference")
                max_new_tokens = gr.Slider(minimum=1, maximum=650, step=1, value=300, label="🪙new tokens in reply")
                top_k = gr.Slider(minimum=1, maximum=100, step=1, value=50, label="🗣️Top K")
                top_p = gr.Slider(minimum=0.1, maximum=1.0, step=0.01, value=0.95, label="🤫Repetition Penalty")
        prompt = gr.Textbox(label="Enter your prompt")
        submit_button.click(fn=generate_text, inputs=[prompt, max_new_tokens, top_k, top_p, do_sample], outputs=output)

    demo.launch()

if __name__ == "__main__":
    main()