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

BASE_ID = "unsloth/DeepSeek-R1-Distill-Llama-8B-unsloth-bnb-4bit"
ADAPTER_ID = "jhenberthf/marites-ai"

print(f"Loading base {BASE_ID} on CPU (fp32)...", flush=True)
tokenizer = AutoTokenizer.from_pretrained(ADAPTER_ID)
base = AutoModelForCausalLM.from_pretrained(
    BASE_ID, torch_dtype=torch.float32, low_cpu_mem_usage=True
)
print(f"Loading LoRA adapter {ADAPTER_ID}...", flush=True)
model = PeftModel.from_pretrained(base, ADAPTER_ID)
model = model.to("cpu")
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
print("Model loaded.", flush=True)

SYSTEM_PROMPT = "You are Maritis AI, a helpful assistant fluent in Filipino and Philippine languages."


def respond(message, history):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for user_msg, bot_msg in history:
        if bot_msg is not None:
            messages.append({"role": "user", "content": user_msg})
            messages.append({"role": "assistant", "content": bot_msg})
    messages.append({"role": "user", "content": message})

    outputs = pipe(
        messages,
        max_new_tokens=512,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1,
    )
    result = outputs[0]["generated_text"]
    if isinstance(result, list):
        return result[-1]["content"]
    return str(result)


demo = gr.ChatInterface(
    respond,
    title="🚀 Maritis AI",
    description=(
        "Local chat assistant using **jhenberthf/marites-ai** "
        "(DeepSeek-R1-Distill-Qwen-1.5B + Filipino/PH-languages LoRA), running on CPU."
    ),
    examples=[
        "Kumusta? Kwentuhan tayo.",
        "Explain what a neural network is in simple Tagalog.",
        "Sumulat ka ng maikling tula tungkol sa AI.",
    ],
    cache_examples=False,
    theme=gr.themes.Soft(),
)

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    demo.launch(server_name="0.0.0.0", server_port=port)