Spaces:
Sleeping
Sleeping
File size: 3,682 Bytes
640b9e7 | 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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | import os
import json
import dotenv
from openai import OpenAI
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.document_loaders import DirectoryLoader
from utils import DOC_CHUNK_SIZE, DOC_CHUNK_OVERLAP, DOC_DIRECTORY, EMBEDDING_FILE
'''
This file handles the loading and embedding of documents.
Saves the embeddings to a json file specified by EMBEDDING_FILE.
Uses OpenAI's text-embedding-3-small model for embeddings.
Supported formats: pdf
Format of saved json file:
List of dictionaries with keys:
'doc_id': {filename}_{chunk number}
'embeddings': List of embeddings for the chunk, one embedding for each character in chunk
Specifications:
OpenAI embeds with a dimension of 1536 per character
'''
# Load environment variables
dotenv.load_dotenv()
def load_documents():
'''
Load documents from a specified directory into a list
Returns:
documents: List of documents loaded from the directory, split by chunks
'''
# Create document loaders
pdf_loader = DirectoryLoader(DOC_DIRECTORY, glob='*.pdf')
docx_loader = DirectoryLoader(DOC_DIRECTORY, glob='*.docx')
txt_loader = DirectoryLoader(DOC_DIRECTORY, glob='*.txt')
loaders = [pdf_loader]
# Load documents
print("Loading documents...")
documents = []
for loader in loaders:
try:
documents.extend(loader.load())
except Exception as e:
print(f"Error loading documents: {e}")
if (len(documents) == 0):
print("No documents loaded.")
return []
# Split documents into chunks
text_splitter = CharacterTextSplitter(chunk_size=DOC_CHUNK_SIZE, chunk_overlap=DOC_CHUNK_OVERLAP)
documents = text_splitter.split_documents(documents)
# Iterate to edit metadata to include chunk number
# format = {filename}_{chunk number}
chunk_num = 1
prev_doc_id = documents[0].metadata['source']
for chunk in documents:
print(chunk.metadata)
if chunk.metadata['source'] != prev_doc_id:
chunk_num = 1
prev_doc_id = chunk.metadata['source']
chunk.metadata['source'] = f"{prev_doc_id}_{chunk_num}"
chunk_num += 1
return documents
def embed_documents(documents):
'''
Embed documents using OpenAIEmbeddings
Args:
documents: List of documents to embed
Returns:
List of JSON objects {doc_id, embeddings, metadata}
'''
# Use OpenAI to embed documents
client = OpenAI(
api_key=os.getenv("OPENAI_API_KEY")
)
embeddings = []
print("Embedding documents...")
# Embed each chunk
for chunk in documents:
chunk_embeddings = client.embeddings.create(
model="text-embedding-3-small",
input=chunk.page_content
)
# Extract embeddings from response
chunk_embedding = [record.embedding for record in chunk_embeddings.data]
embeddings.append({
'doc_id': chunk.metadata['source'],
'embeddings': chunk_embedding[0],
'metadata': {'source': chunk.metadata['source'], 'text': chunk.page_content}
})
return embeddings
def save_embeddings(embeddings, filename):
'''
Save generated embedding to a json file
Args:
embeddings: List of embeddings to save
filename: Name of the file to save the embeddings
'''
print("Saving embedding...")
with open(filename, 'w') as file:
json.dump(embeddings, file)
### Main code
documents = load_documents()
if (len(documents) == 0):
exit()
embeddings = embed_documents(documents)
save_embeddings(embeddings, EMBEDDING_FILE) |