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

BASE_MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
LORA_MODEL = "Logi6023/LogiAI"

print("🤖 LogiAI betöltése CPU-n...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float32,
    device_map="cpu",
)
model = PeftModel.from_pretrained(base_model, LORA_MODEL)
model.eval()
print("✅ LogiAI kész!")

TEMPLATE = "### Instruction:\n{}\n\n### Input:\n\n\n### Response:\n"

def logiai_chat(message, history):
    prompt = TEMPLATE.format(message)
    inputs = tokenizer([prompt], return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            temperature=0.7,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return result.split("### Response:")[-1].strip()

demo = gr.ChatInterface(
    fn=logiai_chat,
    title="🤖 LogiAI",
    description="Saját AI asszisztens — kód, fájlok, CMD parancsok (CPU mód)",
    examples=[
        "Ki vagy te?",
        "Írj Python hello world programot!",
        "Írj CMD parancsot ami kiírja az IP címet",
    ],
    theme=gr.themes.Soft(),
)

demo.launch()