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

BASE_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
LORA_MODEL = "Logi6023/LogiAI-tiny"

print("🤖 LogiAI betöltése...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float32,
    device_map="cpu",
)
model = PeftModel.from_pretrained(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",
    examples=["Ki vagy te?", "Írj Python hello world-öt!", "Mit tudsz csinálni?"],
)

demo.launch()