from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

# Замените 'your_model_path' на путь к вашей модели
model_name = "Reflection-70B"

# Настройки квантовки с float32 offload на CPU
bnb_config = BitsAndBytesConfig(
    load_in_8bit=True  # Загрузка в 8-битном формате
)

# Загрузка токенизатора
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Явное распределение устройства и offload float32 на CPU
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,  # Используем объект конфигурации BitsAndBytes
    device_map="auto",  # Автоматическое распределение между CPU и GPU
    load_in_8bit_fp32_cpu_offload=True,  # Offload float32 модулей на CPU
    trust_remote_code=True
)

# Функция для генерации ответа
def generate_response(prompt, max_length=100):
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    with torch.no_grad():
        outputs = model.generate(**inputs, max_length=max_length, pad_token_id=tokenizer.eos_token_id)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Основная функция чата
def chat():
    print("Чат с моделью начат. Для выхода введите 'exit'.")
    while True:
        # Получаем ввод от пользователя
        user_input = input("Вы: ")
        if user_input.lower() == "exit":
            print("Чат завершен.")
            break

        # Генерация ответа
        response = generate_response(user_input)
        print(f"Модель: {response}")

# Запуск чата
chat()