"""Inference with LoRA adapter for ORTOS bot.
Loads base model (pre-quantized 4-bit) + LoRA weights and answers questions.
"""
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

MODEL_NAME = "unsloth/Meta-Llama-3.1-8B-bnb-4bit"
LORA_PATH = "lora_ortos"


def load_lora_model():
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        device_map="auto",
        dtype=torch.bfloat16,
    )
    model = PeftModel.from_pretrained(model, LORA_PATH)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    tokenizer.pad_token = tokenizer.eos_token
    return model, tokenizer


def ask(model, tokenizer, question: str) -> str:
    prompt = f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
    inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
    outputs = model.generate(**inputs, max_new_tokens=256, temperature=0.3)
    response = tokenizer.decode(outputs[0], skip_special_tokens=False)
    parts = response.split("<|start_header_id|>assistant<|end_header_id|>")
    if len(parts) > 1:
        return parts[-1].replace("<|eot_id|>", "").strip()
    return response


if __name__ == "__main__":
    model, tokenizer = load_lora_model()
    print("LoRA модель загружена. Пиши вопрос (или 'exit'):")
    while True:
        q = input("> ")
        if q.lower() == "exit":
            break
        print(ask(model, tokenizer, q))
        print()