import gradio as gr
from llama_cpp import Llama

# 1. Model ko load karein (Yeh sirf ek baar run hoga jab Space start hoga)
print("Loading model... Please wait, isme 1-2 minute lag sakte hain.")
llm = Llama.from_pretrained(
    repo_id="mradermacher/Qwable-9B-Claude-Fable-5-i1-GGUF",
    filename="Qwable-9B-Claude-Fable-5.i1-IQ1_S.gguf",
    n_ctx=2048,    # Context length
    n_threads=4,   # HF Free CPU ke liye 4 threads sahi kaam karenge
)
print("Model loaded successfully!")

# 2. Chat function jo har message par call hoga
def predict(message, history):
    # Chat history ko format karne ke liye (ChatML format)
    prompt = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
    
    for user_msg, bot_msg in history:
        prompt += f"<|im_start|>user\n{user_msg}<|im_end|>\n<|im_start|>assistant\n{bot_msg}<|im_end|>\n"
        
    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
    
    # Text generate karna
    output = llm(
        prompt,
        max_tokens=512,
        stop=["<|im_end|>", "<|im_start|>"],
        echo=False
    )
    
    return output['choices'][0]['text']

# 3. Gradio ChatInterface launch karna
demo = gr.ChatInterface(
    fn=predict,
    title="Qwable 9B Claude Fable 5 (GGUF CPU)",
    description="Running heavily quantized 9B model on Hugging Face Free CPU Space."
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)