import gradio as gr
from transformers import pipeline

print("Loading model...")
# device_map="auto" automatically loads the model onto the GPU if available
pipe = pipeline(
    "image-text-to-text",
    model="huihui-ai/Huihui-Qwen3.5-2B-abliterated",
    device_map="auto"
)
print("Model loaded successfully!")

def respond(message, history):
    # Reconstruct the message history in the format expected by the model
    messages =[]
    
    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({
                "role": "user",
                "content":[{"type": "text", "text": user_msg}]
            })
        if assistant_msg:
            messages.append({
                "role": "assistant",
                "content":[{"type": "text", "text": assistant_msg}]
            })
            
    # Add the current user message
    messages.append({
        "role": "user",
        "content": [{"type": "text", "text": message}]
    })
    
    # Generate the response
    # FIX: max_new_tokens=2048 prevents the output from cutting off at 250 characters
    response = pipe(text=messages, max_new_tokens=2048)
    
    # Extract the assistant's response dynamically
    generated_messages = response[0]['generated_text']
    assistant_reply = generated_messages[-1]['content']
    
    # Return the text properly whether the model outputs a list of dicts or a raw string
    if isinstance(assistant_reply, list):
        reply_text = "".join([item.get("text", "") for item in assistant_reply if item.get("type") == "text"])
        return reply_text
    else:
        return assistant_reply

# FIX: Removed `theme="soft"` which was causing the TypeError crash in Gradio 4
demo = gr.ChatInterface(
    fn=respond,
    title="Huihui Qwen3.5 Chatbot",
    description="A conversational AI using the `huihui-ai/Huihui-Qwen3.5-2B-abliterated` model.",
    examples=["How to make a book?", "Explain quantum physics to a child.", "Write a poem about AI."]
)

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