rodolpho
62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
import chromadb
|
|
from chromadb.utils import embedding_functions
|
|
import json
|
|
import os
|
|
|
|
# Configurações
|
|
CHUNKS_FILE = "fluig_chunks.json"
|
|
DB_PATH = "fluig_vector_db"
|
|
|
|
def init_vector_db():
|
|
# Usar um modelo extremamente leve adequado para o Raspberry Pi
|
|
# 'all-MiniLM-L6-v2' é o padrão ouro para performance vs qualidade
|
|
model_name = "all-MiniLM-L6-v2"
|
|
emb_fn = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=model_name)
|
|
|
|
# Inicializa o cliente persistente (SQLite sob o capô)
|
|
client = chromadb.PersistentClient(path=DB_PATH)
|
|
|
|
# Cria ou obtém a coleção
|
|
collection = client.get_or_create_collection(
|
|
name="fluig_docs",
|
|
embedding_function=emb_fn,
|
|
metadata={"hnsw:space": "cosine"} # Melhor para busca semântica
|
|
)
|
|
|
|
if not os.path.exists(CHUNKS_FILE):
|
|
print(f"Erro: {CHUNKS_FILE} não encontrado. Rode o rag_processor.py primeiro.")
|
|
return
|
|
|
|
with open(CHUNKS_FILE, "r", encoding="utf-8") as f:
|
|
chunks = json.load(f)
|
|
|
|
print(f"Indexando {len(chunks)} chunks no ChromaDB (isso pode levar alguns minutos no Pi)...")
|
|
|
|
ids = []
|
|
documents = []
|
|
metadatas = []
|
|
|
|
for i, chunk in enumerate(chunks):
|
|
ids.append(f"id_{i}")
|
|
documents.append(chunk["content"])
|
|
metadatas.append({
|
|
"title": chunk["metadata"].get("title", ""),
|
|
"source": chunk["metadata"].get("source", ""),
|
|
"path": chunk["metadata"].get("path", "")
|
|
})
|
|
|
|
# Adicionar em lotes para não estourar memória
|
|
batch_size = 100
|
|
for i in range(0, len(documents), batch_size):
|
|
collection.add(
|
|
ids=ids[i:i+batch_size],
|
|
documents=documents[i:i+batch_size],
|
|
metadatas=metadatas[i:i+batch_size]
|
|
)
|
|
print(f"Progresso: {i + len(documents[i:i+batch_size])}/{len(documents)}")
|
|
|
|
print(f"Sucesso! Banco vetorial criado em {DB_PATH}")
|
|
|
|
if __name__ == "__main__":
|
|
init_vector_db()
|