import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import os

# 1. Configuration - Specify the model and the specific GGUF file
model_id = "Jackrong/Qwen3.5-9B-Claude-4.6-Opus-Reasoning-Distilled-v2-GGUF"
filename = "Qwen3.5-9B-Claude-4.6-Opus-Reasoning-Distilled-v2-Q4_K_M.gguf"

print(f"Downloading model: {filename}...")
model_path = hf_hub_download(repo_id=model_id, filename=filename)

# 2. Initialize the model
# n_ctx is the context window. 2048 is a safe starting point for free Hugging Face Spaces.
llm = Llama(model_path=model_path, n_ctx=2048, n_threads=os.cpu_count())

def generate_response(message, history):
    # Construct the prompt for reasoning models
    # They typically look for a specific structure to trigger <think> tags.
    prompt = f"<|im_start|>system\nYou are a helpful assistant with advanced reasoning capabilities. Use <think> tags for your internal logic.<|im_end|>\n"
    
    for user_msg, assistant_msg in history:
        prompt += f"<|im_start|>user\n{user_msg}<|im_end|>\n"
        prompt += f"<|im_start|>assistant\n{assistant_msg}<|im_end|>\n"
    
    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n<think>\n"
    
    output = llm(
        prompt,
        max_tokens=1024,
        stop=["<|im_end|>"],
        stream=True
    )
    
    response = ""
    for chunk in output:
        delta = chunk['choices'][0]['text']
        response += delta
        yield response

# 3. Create the UI
demo = gr.ChatInterface(
    generate_response,
    title="Qwen 3.5 Reasoning Chat (Claude Distilled)",
    description="This Space runs the Qwen3.5-9B reasoning model distilled from Claude 4.6 Opus logic.",
    examples=["Solve the riddle: What has keys but can't open locks?", "Explain quantum entanglement in simple terms."],
)

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