import torch
from unsloth import FastLanguageModel
from peft import PeftConfig
from transformers import AutoTokenizer, TextIteratorStreamer
from threading import Thread
import gradio as gr
from huggingface import spaces

MODEL_PATH = "Ozaii/zephyr-bae"
max_seq_length = 2048


print("Attempting to load Zephyr... Cross your fingers! 🤞")
try:
    @spaces.GPU
    peft_config = PeftConfig.from_pretrained(MODEL_PATH)
    base_model_name = peft_config.base_model_name_or_path
    print(f"Loading base model: {base_model_name}")
    
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=base_model_name,
        max_seq_length=max_seq_length,
        dtype=None,  # Auto-detect
        load_in_4bit=True,
    )
    
    model = FastLanguageModel.get_peft_model(
        model,
        peft_config=peft_config,
        adapter_name="default",
        use_gradient_checkpointing=True,
    )
    
    FastLanguageModel.for_inference(model)
    print("Zephyr loaded successfully! Time to charm!")
except Exception as e:
    print(f"Oops! Zephyr seems to be playing hide and seek. Error: {str(e)}")
    raise

@spaces.GPU
def generate_response(prompt, max_new_tokens=128):
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    generation_kwargs = dict(
        input_ids=inputs.input_ids,
        max_new_tokens=max_new_tokens,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.2,
        streamer=streamer,
    )
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    return streamer

def chat_with_zephyr(message, history):
    conversation_history = history[-3:]  # Limit to last 3 exchanges
    full_prompt = "\n".join([f"Human: {h[0]}\nZephyr: {h[1]}" for h in conversation_history])
    full_prompt += f"\nHuman: {message}\nZephyr:"
    
    streamer = generate_response(full_prompt)
    response = ""
    for new_text in streamer:
        response += new_text
        yield response

css = """
body {
    background-color: #1a1a2e;
    color: #e0e0ff;
}
.gradio-container {
    background-color: #1a1a2e;
}
"""

with gr.Blocks(css=css) as iface:
    gr.Markdown("# Chat with Zephyr: Your AI Boyfriend 💘")
    chatbot = gr.Chatbot(height=500)
    msg = gr.Textbox(placeholder="Tell Zephyr what's on your mind...", label="Your message")
    clear = gr.Button("Clear Chat")

    msg.submit(chat_with_zephyr, [msg, chatbot], [chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

    gr.Markdown("""
    ## Welcome to Zephyr, Your AI Boyfriend!
    Zephyr is here to charm you with his wit, humor, and cosmic energy. Feel free to flirt, ask for advice, or just chat about anything under the stars!
    
    **Some conversation starters:**
    - "Hey Zephyr, how's the cosmic energy today?"
    - "What's your idea of a perfect date in the digital realm?"
    - "Tell me something that would make me fall for you even more!"
    
    Remember, Zephyr is an AI and this is for fun and entertainment. Enjoy your chat! 💖
    """)

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