import gradio as gr
from transformers import pipeline, set_seed
from transformers import pipeline, set_seed, AutoTokenizer, AutoModelForCausalLM

import torch

description = "The models are intended for both research and commercial use in any of the languages included in the training data. The base models are intended either for language generation or to be further fine-tuned for specific use-cases. The instruction-tuned variants can be used as general-purpose assistants, as long as the user is fully aware of the model’s limitations."

joinus = """
## 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/qdfnvSPcqP) On 🤗Huggingface:[MultiTransformer](https://huggingface.co/MultiTransformer) On 🌐Github: [Tonic-AI](https://github.com/tonic-ai) & contribute to🌟 [Build Tonic](https://git.tonic-ai.com/contribute)🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗
"""

device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "BSC-LT/salamandra-2b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")

# Set pad_token_id to eos_token_id for open-end generation
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

def generate_text(prompt, temperature, top_p, max_new_tokens, repetition_penalty):
    # set_seed(42)
    generation_args = {
        "temperature": temperature,
        "top_p": top_p,
        "max_new_tokens": max_new_tokens,
        "repetition_penalty": repetition_penalty,
        "do_sample": True,
        "pad_token_id": tokenizer.eos_token_id
    }
    output = generator(prompt, **generation_args)
    return output[0]["generated_text"]

def update_output(prompt, temperature, top_p, max_new_tokens, repetition_penalty):
    generated_text = generate_text(prompt, temperature, top_p, max_new_tokens, repetition_penalty)
    return generated_text

def update_prompt(example):
    return example

with gr.Blocks() as demo:
    gr.Markdown("# 🙋🏻‍♂️Welcome to Tonic's 📲🦎Salamandra-2b On-Device Demo")
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Group():
                gr.Markdown(description)
        with gr.Column(scale=1):
            with gr.Group():
                gr.Markdown(joinus)

    with gr.Row(): 
        with gr.Column(scale=1):    
            prompt = gr.Textbox(lines=5, label="🙋🏻‍♂️Input Prompt")
            generate_button = gr.Button("Try 📲🦎Salamandra-2b")
            with gr.Accordion("🧪Parameters", open=False):           
                # with gr.Column(scale=1):
                temperature = gr.Slider(0.0, 1.0, value=0.7, label="🌡️Temperature")
                top_p = gr.Slider(0.0, 1.0, value=0.95, label="⚛️Top P")
                max_new_tokens = gr.Slider(1, 4096, value=350, step=1, label="🤑Max New Tokens")
                repetition_penalty = gr.Slider(1.0, 2.0, value=1.2, label="🦜Repetition Penalty")
        
        with gr.Column(scale=1): 
            output = gr.Textbox(lines=10, label=" 📲🦎Salamandra-2b")
        
    generate_button.click(
        update_output,
        inputs=[prompt, temperature, top_p, max_new_tokens, repetition_penalty],
        outputs=output
    )
    
    examples = gr.Examples(
        examples=[
            ["Todo el mundo sabe que vivir en Barcelona es"],
            ["¿Pueblo o ciudad? Una ventaja de vivir en la ciudad es que hay muchas oportunidades de ocio y empleo, así como una gran diversidad de comercios para todos los gustos. Sin embargo, las ciudades suelen ser "],
            ["Llegir ens proporciona"],
            ["What I find more fascinating about languages is that"],
            ["La vie peut être"],
            ["The future of AI is"]
        ],
        inputs=prompt,
        outputs=prompt,
        fn=update_prompt,
        label="Example Prompts"
    )

if __name__ == "__main__":
    demo.launch()