import gradio as gr
from transformers import pipeline

# 1. Setup the Model
# If your custom model is still building, this uses GPT2 as a stable base
model_id = "gpt2" 

try:
    # We use a pipeline for easy text generation
    pipe = pipeline("text-generation", model=model_id)
except Exception as e:
    print(f"Error loading model: {e}")

def zew_chat(message, history):
    # This is the "Pattern" we want ZewAI to follow
    # It teaches the AI: "When you see a Question, give a short Assistant answer"
    training_pattern = (
        "Question: Who are you?\n"
        "Assistant: I am ZewAI, the official assistant for Zako Technology Labs.\n\n"
        "Question: What is ZTL?\n"
        "Assistant: ZTL stands for Zako Technology Labs, a software and hardware innovation group.\n\n"
        "Question: Tell me about ZewpolOS.\n"
        "Assistant: ZewpolOS is an advanced web-based operating system project from ZTL.\n\n"
    )
    
    # We add your actual message to the end of the pattern
    full_prompt = f"{training_pattern}Question: {message}\nAssistant:"
    
    output = pipe(
        full_prompt, 
        max_new_tokens=40,      # Short answers are better for GPT-2
        do_sample=True,          
        temperature=0.3,         # Very low temperature keeps it on track
        top_k=50,
        repetition_penalty=1.5,  # High penalty stops it from repeating "the server"
        pad_token_id=50256
    )
    
    raw_text = output[0]['generated_text']
    
    # This magic line cuts off everything EXCEPT the final Assistant answer
    if "Assistant:" in raw_text:
        # We take the very last Assistant response in the text
        return raw_text.split("Assistant:")[-1].strip().split("\n")[0]
    
    return "I am ZewAI. How can I help with ZTL?"

    # We format the prompt to lead the AI into an 'Assistant' response
    full_prompt = f"{system_instruction}\nUser: {message}\nAssistant:"
    
    # GENERATION SETTINGS: These stop the "random stuff"
    output = pipe(
        full_prompt, 
        max_new_tokens=60,      # Give it enough room to answer
        do_sample=True,          # Allow for creative (but controlled) words
        temperature=0.4,         # Lower = more focused and less random
        top_k=40,                # Only use the best 40 words
        top_p=0.9,               # Filter out rare/weird words
        repetition_penalty=1.3,  # Heavily stop it from repeating the same thing
        pad_token_id=50256       # Required for GPT2
    )
    
    # CLEANING THE TEXT: 
    # This removes the System Instruction and the User's question from the final answer
    raw_text = output[0]['generated_text']
    if "Assistant:" in raw_text:
        return raw_text.split("Assistant:")[-1].strip()
    
    return raw_text

# 2. Build the Interface
# We remove 'theme' because modern Gradio handles it differently
demo = gr.ChatInterface(
    fn=zew_chat, 
    title="ZewAI Official (ZTL)",
    description="Developed by Zako Technology Labs for ZewpolOS 4 pine trunk."
)

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