init_models.py

import torch
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig


from projeto.app import cfg, load_entries


def debug_json_structure(path):
    """Debug para ver a estrutura real do JSON"""
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)

    print("🔍 DEBUG DA ESTRUTURA DO JSON:")

    # Ver as primeiras 3 chaves para ver a estrutura completa
    for i, (key, value) in enumerate(list(data.items())[:3]):
        print(f"\n--- Chave {i + 1}: '{key}' ---")
        print(f"Tipo do valor: {type(value)}")
        if isinstance(value, dict):
            print(f"Campos: {list(value.keys())}")
            for k, v in value.items():
                if k == "embedding":
                    print(f"  {k}: [lista com {len(v) if isinstance(v, list) else '?'} elementos]")
                else:
                    print(f"  {k}: {str(v)[:100]}{'...' if len(str(v)) > 100 else ''}")
        print("---")

    # Procurar especificamente por "Faust" para ver sua estrutura
    print("\n🔍 PROCURANDO POR 'Faust' NO JSON:")
    faust_found = False
    for key, value in data.items():
        if "Faust" in key or (
                isinstance(value, dict) and "Faust" in str(value.get('title', '')) + str(value.get('term', ''))):
            print(f"Encontrado Faust na chave: '{key}'")
            print(f"Estrutura: {value}")
            faust_found = True
            break

    if not faust_found:
        print("Faust não encontrado nas primeiras verificações")

def init_models():
    """Initialize all models and load data"""
    # 1. Load embedding model for semantic search
    embed_model = SentenceTransformer("all-MiniLM-L6-v2", device=cfg.DEVICE)

    # 2. Load tokenizer and model for text generation
    tokenizer = AutoTokenizer.from_pretrained(cfg.MODEL_ID)

    # Add pad token if missing
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # Load model without quantization to avoid bitsandbytes issues
    model = AutoModelForCausalLM.from_pretrained(
        cfg.MODEL_ID,
        device_map="auto",
        torch_dtype=torch.float16 if cfg.DEVICE == "cuda" else torch.float32
    )

    # 3. Load entries and vectors
    print("Carregando embeddings...")
    entries, vectors = load_entries(cfg.EMBEDDINGS_FILE)

    if len(entries) == 0:
        raise Exception("Nenhuma entrada foi carregada! Verifique o arquivo de embeddings.")

    return embed_model, tokenizer, model, entries, vectors