import gradio as gr
import torch
import os
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# Token from Secrets
hf_token = os.environ.get("HF_TOKEN")

model_id = "unsloth/qwen2.5-7b-bnb-4bit"
adapter_id = "Alauddin123/BongoAI-V1.0"

def load_bongo():
    try:
        print("--- Loading Tokenizer ---")
        tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
        
        print("--- Loading 7B Model (CPU Mode) ---")
        model = AutoModelForCausalLM.from_pretrained(
            model_id, 
            token=hf_token, 
            trust_remote_code=True,
            device_map="cpu", 
            low_cpu_mem_usage=True,
            torch_dtype=torch.float32
        )
        
        print("--- Applying Adapter ---")
        model = PeftModel.from_pretrained(model, adapter_id, token=hf_token)
        print("--- SUCCESS: BongoAI is Online! ---")
        return tokenizer, model
    except Exception as e:
        print(f"CRITICAL ERROR: {str(e)}")
        return None, str(e)

tokenizer, bongo_model = load_bongo()

def chat(message, history):
    if tokenizer is None:
        return f"System Error: {bongo_model}"
    
    prompt = f"### Instruction:\n{message}\n\n### Response:\n"
    inputs = tokenizer(prompt, return_tensors="pt")
    
    with torch.no_grad():
        outputs = bongo_model.generate(
            **inputs, 
            max_new_tokens=128,
            temperature=0.7,
            do_sample=True
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    if "### Response:" in response:
        response = response.split("### Response:")[-1].strip()
    
    return response

# Simple interface without tabs for now
demo = gr.ChatInterface(
    fn=chat,
    title="BongoAI 7B",
    description="Ask me anything!"
)

demo.launch(server_name="0.0.0.0", server_port=7860)