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


hf_token = os.getenv("YOUR_HF_TOKEN")

# Load model and tokenizer
print("Loading model and tokenizer...")
model_path = "microsoft/Phi-4-mini-instruct"  # Can be changed to local path "./Phi-4-Mini-Instruct"

tokenizer = AutoTokenizer.from_pretrained(
    model_path, 
    padding_side="left",
    token=hf_token,
    trust_remote_code=True
)

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",
    attn_implementation="flash_attention_2",
    torch_dtype="auto",
    token=hf_token,
    trust_remote_code=True
)

# Create pipeline for easier inference
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
)

print("Model and tokenizer loaded successfully!")

# Format chat history to messages format
def format_chat_history(message, history):
    messages = [
        {"role": "system", "content": "You are a helpful AI assistant."}
    ]
    
    # Add chat history
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    
    # Add current message
    messages.append({"role": "user", "content": message})
    
    return messages

# Streaming response generator
def predict(message, history):
    messages = format_chat_history(message, history)
    
    generation_args = {
        "max_new_tokens": 1024,
        "return_full_text": False,
        "temperature": 0.001,
        "top_p": 1.0,
        "do_sample": True,
        "streamer": None,  # Will be set in the generator
    }
    
    # Initialize an empty response
    partial_message = ""
    history_with_message = history + [[message, partial_message]]
    
    # Create a TextIteratorStreamer for streaming generation
    from transformers import TextIteratorStreamer
    from threading import Thread
    
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    generation_args["streamer"] = streamer
    
    # Start a separate thread for generation
    thread = Thread(target=pipe, args=(messages,), kwargs=generation_args)
    thread.start()
    
    # Stream the response
    for new_text in streamer:
        partial_message += new_text
        yield history + [[message, partial_message]]

# Create the Gradio interface
css = """
.chatbot-container {max-width: 800px; margin: auto;}
.chat-header {text-align: center; margin-bottom: 20px;}
"""

with gr.Blocks(css=css) as demo:
    gr.HTML("<div class='chat-header'><h1>Phi-4 Mini Chatbot</h1></div>")
    
    with gr.Column(elem_classes="chatbot-container"):
        chatbot = gr.Chatbot(height=400)
        msg = gr.Textbox(placeholder="Type your message here...", label="Input")
        clear = gr.Button("Clear Conversation")
        
        msg.submit(predict, [msg, chatbot], [chatbot], queue=True, api_name="chat").then(
            lambda: "", None, [msg]
        )
        clear.click(lambda: None, None, chatbot, queue=False)
    """)

# Launch the app
demo.launch(share=True)  # Set share=False if you don't want a public link