import torch
from peft import PeftModel, PeftConfig
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
from threading import Thread
import gradio as gr

MODEL_PATH = "Ozaii/zephyr-bae"

print("Attempting to load Zephyr... Cross your fingers! 🤞")

try:
    peft_config = PeftConfig.from_pretrained(MODEL_PATH)
    base_model = AutoModelForCausalLM.from_pretrained(
        peft_config.base_model_name_or_path,
        torch_dtype=torch.float16,
        device_map="auto"
    )
    model = PeftModel.from_pretrained(base_model, MODEL_PATH)
    tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path)
    tokenizer.pad_token = tokenizer.eos_token
    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

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()