File size: 3,862 Bytes
8379fcb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
# local_process_embeddings.py
from pathlib import Path
import ollama
import faiss
import numpy as np
import os
import glob
import pickle # For saving documents list
# (Copy your inicializar_rag function here)
# Make sure it's adapted to just return documents and index without Gradio parts
def inicializar_rag():
"""
Inicializa o modelo de embeddings, documentos e índice FAISS a partir dos arquivos de capítulo.
Adaptado para retornar apenas documentos e índice FAISS, sem partes do Gradio.
"""
# Diretório onde os arquivos de capítulo estão localizados
base_dir = Path(__file__).resolve().parent
diretorio_rag = base_dir / "RAG"
print(f">>> Looking for files under: {diretorio_rag}")
if not diretorio_rag.exists():
raise FileNotFoundError(
f"O diretório {diretorio_rag} não existe. Verifique se ele foi copiado corretamente para o container."
)
# Use pathlib's glob for more consistent path handling
arquivos = list(diretorio_rag.glob("[bsde][1-9].txt"))
caminhos_arquivos = [str(arquivo) for arquivo in arquivos] # convert pathlib objects to strings
documentos = []
# Carrega e processa cada arquivo de capítulo
for caminho in caminhos_arquivos:
with open(caminho, "r", encoding="utf-8") as f:
conteudo = f.read().split("\n\n")
documentos.extend(conteudo)
if not documentos:
raise ValueError("Nenhum documento foi carregado. Verifique os caminhos dos arquivos.")
print(f"Trabalhando em {len(documentos)} documentos ...")
# Gera embeddings usando o Ollama
embeddings = []
for index, doc in enumerate(documentos):
if index % 200 == 0:
print(f"= = = {index} Embeddings... = = = doc: {doc[:20]}") # added print of the doc
try:
response = ollama.embed(model="nomic-embed-text", input=doc)
embeddings.append(response["embeddings"][0]) # Access the first element of the embeddings list.
except Exception as e:
print(f"Error generating embedding for document {index}: {e}")
# Handle the error, e.g., skip this document or raise an exception
embeddings.append(None) # Append None to keep the lists aligned
continue
# Remove any None embeddings (documents that failed to embed)
embeddings_valid = [emb for emb in embeddings if emb is not None]
documentos_valid = [doc for i, doc in enumerate(documentos) if embeddings[i] is not None]
if not embeddings_valid:
raise ValueError("No valid embeddings were generated.")
print(f"Trabalhando em {len(embeddings_valid)} embeddings...")
embeddings_np = np.array(embeddings_valid, dtype=np.float32)
# Cria o índice FAISS
dimension = embeddings_np.shape[1]
faiss_index = faiss.IndexFlatL2(dimension)
faiss_index.add(embeddings_np)
return documentos_valid, faiss_index
if __name__ == "__main__":
print("Inicializando RAG localmente...")
# Ensure the RAG directory is correctly found relative to this script
# or use an absolute path for local execution.
# For example, if RAG is in the same dir as this script:
# base_dir = Path(__file__).resolve().parent
# diretorio_rag_local = base_dir / "RAG"
documentos_global, index_global = inicializar_rag() # Call your function
print(f"Total de documentos processados: {len(documentos_global)}")
print(f"Dimensão do índice FAISS: {index_global.d}, Total de vetores: {index_global.ntotal}")
# Save the documents list
with open("documentos.pkl", "wb") as f_docs:
pickle.dump(documentos_global, f_docs)
print("Lista de documentos salva em documentos.pkl")
# Save the FAISS index
faiss.write_index(index_global, "faiss_index.idx")
print("Índice FAISS salvo em faiss_index.idx")
|