import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import os
import re
import spaces

# The model to call
model_id = "Atlas-labs/mini-fable-5-qwen-merged"
token = os.getenv("HF_TOKEN")

print(f"Loading tokenizer and model...")
tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    torch_dtype=torch.float16, 
    device_map="auto",
    token=token
)

def parse_fable_response(raw_text, user_input):
    thought = ""
    thought_match = re.search(r'<thought>(.*?)</thought>', raw_text, re.DOTALL | re.IGNORECASE)
    if thought_match:
        thought = thought_match.group(1).strip()
    
    answer = raw_text
    if "</thought>" in raw_text.lower():
        answer = re.split(r'</thought>', raw_text, flags=re.IGNORECASE)[-1]
    
    answer = re.sub(r'</?thought>', '', answer, flags=re.IGNORECASE)
    answer = re.sub(r'</?response>', '', answer, flags=re.IGNORECASE)
    answer = answer.replace(user_input, "").strip()

    if thought:
        return f"💡 **Reasoning:**\n> *{thought}*\n\n{answer}"
    return answer

@spaces.GPU
def chat(message, history):
    prompt = f"<thought>\nAnalyzing request: {message}\n"
    inputs = tokenizer(prompt + message, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs, 
            max_new_tokens=512, 
            temperature=0.7, 
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    
    raw_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return parse_fable_response(raw_response, message)

demo = gr.ChatInterface(
    fn=chat, 
    title="Mini Fable 5 (ZeroGPU)", 
    description="High-speed reasoning engine by Atlas Labs. Powered by Hugging Face ZeroGPU.",
    examples=["Explain the theory of relativity.", "Write a Python script for a binary search.", "Why is the sky blue?"]
)

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