import gradio as gr
import os
import spaces
from transformers import GemmaTokenizer, AutoModelForCausalLM

# Set an environment variable
HF_TOKEN = os.environ.get("HF_TOKEN", None)

# Load the tokenizer and model
tokenizer = GemmaTokenizer.from_pretrained("google/codegemma-7b-it")
model = AutoModelForCausalLM.from_pretrained("google/codegemma-7b-it").to("cuda:0")


def codegemma(message: str, history: list, temperature: float, max_new_tokens: int) -> str:
    """
    Generate a response using the CodeGemma model.

    Args:
        message (str): The input message.
        history (list): The conversation history used by ChatInterface.
        temperature (float): The temperature for generating the response.
        max_new_tokens (int): The maximum number of new tokens to generate.

    Returns:
        str: The generated response.
    """
    input_ids = tokenizer(message, return_tensors="pt")
    outputs = model.generate(
        **input_ids,
        temperature=temperature,
        max_new_tokens=max_new_tokens,
    )
    response = tokenizer.decode(outputs[0])
    return response

    
placeholder = """
<img src="https://huggingface.co/spaces/ysharma/CodeGemma/resolve/main/gemma_lockup_vertical_full-color_rgb.png" style="width:40%">
<b>CodeGemma-7B-IT</b>
"""

# Gradio block
with gr.Blocks(fill_height=True) as demo:
    gr.Markdown("# GEMMA-7b-IT")
    #with gr.Tab('CodeGemma Chatbot'):
    gr.ChatInterface(codegemma, 
                   examples=[["Write a Python function to calculate the nth fibonacci number."]], 
                   fill_height=True,
                   additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
                   additional_inputs=[
                       gr.Slider(0, 1, 0.95, label="Temperature", render=False),
                       gr.Slider(128, 4096, 512, label="Max new tokens", render=False ),
                       ],
                   )
    

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