import gradio as gr
from huggingface_hub import InferenceClient
import os
from dotenv import load_dotenv

load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")

model_list = ["google/gemma-2-2b-it", "google/gemma-2-9b-it", "google/gemma-2-27b-it"]

def respond(
    message,
    history: list[tuple[str, str]],
    model_id,
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    client = InferenceClient(
        model_id,
        token=HF_TOKEN,
    )
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content

        response += token
        yield response

gemma_chatbot = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Dropdown(
            choices=model_list,
            label="Model",
            value="google/gemma-2-27b-it",
        ),
        gr.Textbox(
            value="You are a friendly Chatbot.",
            label="System message"
        ),
        gr.Slider(
            minimum=1,
            maximum=4096,
            value=512,
            step=1,
            label="Max new tokens"
        ),
        gr.Slider(
            minimum=0.1,
            maximum=4.0,
            value=0.7,
            step=0.1,
            label="Temperature"
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)