import gradio # For the GUI
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline # For LLMs!

# Instantiate the model that we'll be calling. This is a tiny one!
MODEL_ID = "HuggingFaceTB/SmolLM2-135M-Instruct"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
pipe = pipeline(
    task="text-generation",
    model=AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
    ),
    tokenizer=tokenizer
)

# Create a helper to format the chat message appropriately
def _format_chat(system_prompt: str, user_prompt: str, history: list[dict]) -> str:
    """Format messages using chat template if available."""
    messages = [{"role": "system", "content": system_prompt}]
    messages.extend(history)
    messages.append({"role": "user", "content": user_prompt})

    template = getattr(tokenizer, "chat_template", None)
    return tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

# Create an LLM helper function
def _llm_generate(prompt: str, max_tokens: int, temperature: float) -> str:
    out = pipe(
        prompt,
        max_new_tokens=max_tokens,
        do_sample=True,
        temperature=temperature,
        return_full_text=False
    )
    return out[0]["generated_text"].strip()

# Wrap the above in a response function that will be the primary entry point for the GUI
def respond(
    message,
    history: list[dict[str, str]],
    system_message,
    max_tokens,
    temperature,
):

    formatted = _format_chat(system_message, message, history)
    response = _llm_generate(
        formatted,
        max_tokens=max_tokens,
        temperature=temperature,
    )
    return response

# Last but not least, here's the UI! This one is pretty simple.
chatbot = gradio.ChatInterface(
    respond,
    type="messages",
    additional_inputs=[
        gradio.Textbox(
            value=(
                "You are a chatbot that helps come up with fancy, exciting names for research papers."
            ),
            label="System message"
        ),
        gradio.Slider(minimum=0, maximum=256, value=128, step=16, label="Max new tokens"),
        gradio.Slider(minimum=0.1, maximum=2.0, value=0.9, step=0.1, label="Temperature"),
    ],
)

with gradio.Blocks() as demo:


    # Let's start by adding a title and introduction
    gradio.Markdown(
        "# A Simple Chatbot"
    )
    gradio.Markdown(
        "This app runs a simple chatbot on CPU and exposes a few backen parameters too."
    )
    
    # Now insert the chat window
    chatbot.render()

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