from llama_cpp import Llama
from huggingface_hub import hf_hub_download
import os

# Параметры модели (Q5_K_M — хороший баланс скорость/качество на CPU)
REPO_ID = "bartowski/Qwen2.5-3B-Instruct-GGUF"
FILENAME = "Qwen2.5-3B-Instruct-Q5_K_M.gguf"  # если хочешь другой квант — Q4_K_M (быстрее), Q6_K (лучше качество)

# Директория кэша (в HF Spaces /tmp сохраняется между перезапусками, модель не скачивается заново)
CACHE_DIR = "/tmp/hf/models"
os.makedirs(CACHE_DIR, exist_ok=True)

llm = None


def load_model():
    global llm

    if llm is not None:
        return llm  # уже загружено — не тратим время

    print(f"Скачиваем/находим модель {REPO_ID}/{FILENAME} ... (первый раз ~5–15 мин на HF CPU)")

    # Скачивание (или берём из кэша, если уже есть)
    model_path = hf_hub_download(
        repo_id=REPO_ID,
        filename=FILENAME,
        cache_dir=CACHE_DIR,
        local_dir=CACHE_DIR,  # сохраняем явно
        local_dir_use_symlinks=False  # без ссылок, чтобы файл был физически
    )

    print(f"Путь к модели: {model_path}")

    llm = Llama(
        model_path=model_path,
        n_ctx=8192,  # контекст — достаточно для длинных описаний + промпта
        n_threads=0,  # все доступные ядра CPU
        n_gpu_layers=0,  # строго CPU
        n_batch=512,
        verbose=True  # логи в консоль Spaces — увидишь токены/сек
    )

    print("Модель Qwen2.5-7B-Instruct успешно загружена в CPU-режиме")
    return llm