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

title = "# 🙋🏻‍♂️Tonic's ✒️InkubaLM-0.4B"
description = """✒️InkubaLM has been trained from scratch using 1.9 billion tokens of data for five African languages, along with English and French data, totaling 2.4 billion tokens of data.  It is capable of understanding and generating content in five African languages: Swahili, Yoruba, Hausa, isiZulu, and isiXhosa, as well as English and French. 
### 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:[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 🤗
"""

hf_token = os.getenv("HF_TOKEN")
# Load the model and tokenizer
model_name = "lelapa/InkubaLM-0.4B"
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, token=hf_token)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, token=hf_token)

# Move model to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

def generate_text(prompt, max_length, repetition_penalty, temperature):
    # Tokenize the input and create attention mask
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    input_ids = inputs.input_ids
    attention_mask = inputs.attention_mask
    
    # Generate the text using the model, with the attention mask and temperature
    outputs = model.generate(
        input_ids, 
        attention_mask=attention_mask,
        max_length=max_length, 
        repetition_penalty=repetition_penalty, 
        temperature=temperature,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=True
    )
    
    # Decode the generated tokens and return the result
    generated_text = tokenizer.batch_decode(outputs[:, input_ids.shape[1]:-1], skip_special_tokens=True)[0].strip()
    return generated_text

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown(title)
    gr.Markdown(description)
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Enter your prompt here:", placeholder="Today I planned to ...")
            max_length = gr.Slider(label="Max Length", minimum=70, maximum=1000, step=50, value=200)
            repetition_penalty = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, step=0.1, value=1.2)
            temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, step=0.1, value=0.5)  # Added slider for temperature
            submit_button = gr.Button("Generate")
        
        with gr.Column():
            output = gr.Textbox(label="✒️Inkuba.4B:")

    submit_button.click(generate_text, inputs=[prompt, max_length, repetition_penalty, temperature], outputs=output)

# Launch the demo
demo.launch()