import os
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig

# Load config directly from Qwen3-4B base (bypasses broken config.json)
config = AutoConfig.from_pretrained("Qwen/Qwen3-4B", trust_remote_code=True)

# Load your finetuned weights with that config
model = AutoModelForCausalLM.from_pretrained(
    "Weblake/Lake-0.2",
    config=config,
    trust_remote_code=True,
    device_map="auto",
    ignore_mismatched_sizes=True
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B", trust_remote_code=True)

SYSTEM_PROMPT = "You are Lake, an AI assistant made by Weblake Inc. You are not ChatGPT."

def chat(message, history):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for h in history:
        messages.append({"role": "user", "content": h[0]})
        messages.append({"role": "assistant", "content": h[1]})
    messages.append({"role": "user", "content": message})

    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=512)
    return tokenizer.decode(outputs[0][inputs.input_ids.shape[-1]:], skip_special_tokens=True)

gr.ChatInterface(chat).launch()