# SmolLM2-1.7B-Instruct is no longer served by any Hugging Face Inference Provider,
# so this Space now runs the model itself on CPU with transformers.
from threading import Thread

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

MODEL_ID = "HuggingFaceTB/SmolLM2-1.7B-Instruct"

torch.set_num_threads(max(1, torch.get_num_threads()))
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32)
model.eval()


def content_to_text(content):
    """Gradio 6 message content may be a string or a list of parts; flatten to text."""
    if isinstance(content, str):
        return content
    if isinstance(content, dict):
        return content.get("text") or ""
    if isinstance(content, (list, tuple)):
        return "".join(content_to_text(p) for p in content)
    return "" if content is None else str(content)


def respond(
    message,
    history,
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = []
    if system_message:
        messages.append({"role": "system", "content": system_message})

    for msg in history or []:
        if isinstance(msg, dict) and msg.get("role") in ("user", "assistant"):
            text = content_to_text(msg.get("content"))
            if text:
                messages.append({"role": msg["role"], "content": text})

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

    inputs = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
    )
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=int(max_tokens),
        do_sample=temperature > 0,
        temperature=float(temperature),
        top_p=float(top_p),
    )
    thread = Thread(target=model.generate, kwargs=generate_kwargs)
    thread.start()

    response = ""
    for token in streamer:
        response += token
        yield response
    thread.join()

chatbot = gr.Chatbot(height=600)

demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=256, 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",
        ),

    ],
    fill_height=True,
    chatbot=chatbot,
)
if __name__ == "__main__":
    demo.launch(theme="Nymbo/Nymbo_Theme")