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

# Initialize model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    "microsoft/Phi-3.5-mini-instruct",
    trust_remote_code=True,
    torch_dtype=torch.float32,  # Force float32 for CPU compatibility
)

tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3.5-mini-instruct")

# Create pipeline
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=-1,  # Force CPU usage (-1 = CPU)
)

def respond(message, chat_history):
    # Format the chat history with system prompt
    system_prompt = "You are a helpful AI assistant."
    messages = [{"role": "system", "content": system_prompt}]
    
    # Add previous chat history
    for user_msg, bot_msg in chat_history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    
    # Add new user message
    messages.append({"role": "user", "content": message})
    
    # Format for the model
    formatted_input = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    
    # Generate response
    generation_args = {
        "max_new_tokens": 200,  # Reduced for CPU performance
        "temperature": 0.7,
        "do_sample": True,
    }
    
    output = pipe(formatted_input, **generation_args)
    response = output[0]['generated_text']
    
    # Remove the input text from the response
    if response.startswith(formatted_input):
        response = response[len(formatted_input):]
    
    return response

# Create Gradio interface
demo = gr.ChatInterface(
    respond,
    chatbot=gr.Chatbot(height=400),
    textbox=gr.Textbox(placeholder="Ask me anything...", container=False, scale=7),
    title="Phi-3 Mini Chat (CPU)",
    examples=["What's 2x + 3 = 7?", "How to make banana smoothie?", "Explain quantum computing simply"],
    cache_examples=False,
)

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