import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel

# Загрузка модели
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

tokenizer = AutoTokenizer.from_pretrained(
    "unsloth/qwen2.5-7b-instruct-unsloth-bnb-4bit",
    trust_remote_code=True,
)

model = AutoModelForCausalLM.from_pretrained(
    "unsloth/qwen2.5-7b-instruct-unsloth-bnb-4bit",
    device_map="auto",
    quantization_config=bnb_config,
    trust_remote_code=True,
)

# Загрузка LoRA адаптера
model = PeftModel.from_pretrained(model, "Muzurixkurumi/amalia-qwen2.5-7b-lora")
model = model.merge_and_unload()
model.eval()

# Загрузка конфигурации
import json
with open("inference_config.json", "r", encoding="utf-8") as f:
    config = json.load(f)

SYSTEM_PROMPT = config["system_prompt"]
FEW_SHOT_EXAMPLES = config["few_shot_examples"]
MAX_HISTORY_PAIRS = config["inference_config"]["max_history_pairs"]
gen_params = config["generation_config"]

# Генерация ответа
conversation_history = []  # [(user_msg, assistant_msg), ...]

def generate_response(user_message):
    # Формируем промпт
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    
    # Few-shot примеры
    for example in FEW_SHOT_EXAMPLES:
        messages.append({"role": "user", "content": example["user"]})
        messages.append({"role": "assistant", "content": example["assistant"]})
    
    # Rolling window истории (КРИТИЧНО!)
    for user_msg, assistant_msg in conversation_history[-MAX_HISTORY_PAIRS:]:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    
    # Текущий запрос
    messages.append({"role": "user", "content": user_message})
    
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=gen_params["max_new_tokens"],
            temperature=gen_params["temperature"],
            top_p=gen_params["top_p"],
            top_k=gen_params["top_k"],
            do_sample=gen_params["do_sample"],
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    
    response = tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[1]:], 
        skip_special_tokens=True
    ).strip()
    
    # Сохраняем в историю
    conversation_history.append((user_message, response))
    
    return response

# Пример использования
print(generate_response("привет амалия!"))
# Output: йо

print(generate_response("как дела?"))
# Output: норм, а у тебя?

⚠️ ВАЖНО: Rolling Window

ОБЯЗАТЕЛЬНО используйте rolling window (max_history_pairs=5)!

Без ограничения истории модель начинает переключаться на китайский при контексте >700 токенов (это особенность базовой Qwen модели).

🎛️ Параметры генерации

Оптимальные параметры (найдены grid search):

python

temperature = 0.7
top_p = 0.9
top_k = 50  # НЕ УВЕЛИЧИВАТЬ! >50 → риск китайского
max_new_tokens = 120

Не рекомендуется:

temperature > 0.8 (китайский + артефакты)
top_k > 50 (китайский появляется раньше)
Полная история без rolling window (китайский после 7-8 реплик)

📈 Результаты тестирования

Короткие диалоги (1-3 реплики): ✅ Идеально Длинные диалоги (16+ реплик с rolling window): ✅ Идеально Без rolling window: ❌ Китайский после 7 реплик

🛠️ Технические требования

GPU: 8+ GB VRAM (для 4-bit quantization)
RAM: 16+ GB
Python: 3.8+

📝 Лицензия

Apache 2.0 (следует лицензии базовой модели Qwen)

🙏 Благодарности

Qwen Team за базовую модель
Unsloth за оптимизированную версию