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

# Základní model (3B varianta LLaMA 3.2)
base_model_name = "unsloth/Llama-3.2-3B"
tokenizer = AutoTokenizer.from_pretrained(base_model_name)
model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    device_map="auto",           # důležité pro GPU
    torch_dtype=torch.float16,   # nebo bfloat16, pokud máš A100
)

# Načtení LoRA adaptéru
lora_model_path = "Artaccann/OSOL_backend"  # <- správně s malým "b", repozitář Hugging Face je case-sensitive!
model = PeftModel.from_pretrained(
    model,
    lora_model_path,
    device_map="auto",
    torch_dtype=torch.float16,
)

model.eval()  # vypne trénovací režim

# Jednoduchá funkce pro odpověď
def chat(user_input):
    prompt = f"{user_input}"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=250)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# UI rozhraní
gr.Interface(
    fn=chat,
    inputs="text",
    outputs="text",
    title="OSOL AI Chatbot",
    description="Mluvíš s postavou z OSOL, která používá LoRA přizpůsobenou pro tento projekt."
)