Spaces:
Sleeping
Sleeping
File size: 8,747 Bytes
c6907ac |
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 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 |
from utilities.constants import (
CHUNKING_STRATEGY_TABLE_AWARE,
CHUNKING_STRATEGY_SECTION_BASED,
CHUNKING_STRATEGY_SEMANTIC
)
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.docstore.document import Document
from langchain_community.vectorstores import Qdrant
from langchain_openai.embeddings import OpenAIEmbeddings
import numpy as np
import pdfplumber
import re
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import tiktoken
from utilities.debugger import dprint
def create_vector_store(app_state, model_run_state, **kwargs):
for key, value in kwargs.items():
if hasattr(model_run_state, key):
setattr(model_run_state, key, value)
else:
print(f"Warning: {key} is not an attribute of the state object")
# Rest of your create_vector_store logic
dprint(app_state, f"Chunk size after update: {model_run_state.chunk_size}")
create_chunked_documents(app_state, model_run_state)
qdrant_vectorstore = Qdrant.from_documents(
documents=model_run_state.combined_document_objects,
embedding=model_run_state.embedding_model,
location=":memory:"
)
qdrant_retriever = qdrant_vectorstore.as_retriever()
model_run_state.retriever = qdrant_retriever
print("Vector store created")
def tiktoken_len(text):
tokens = tiktoken.encoding_for_model("gpt-4o").encode(
text,
)
return len(tokens)
def create_chunked_documents(app_state, model_run_state):
dprint(app_state, model_run_state.chunking_strategy)
if model_run_state.chunking_strategy == CHUNKING_STRATEGY_TABLE_AWARE:
combined_document_objects = chunk_with_table_aware(app_state, model_run_state)
elif model_run_state.chunking_strategy == CHUNKING_STRATEGY_SECTION_BASED:
combined_document_objects = chunk_with_section_based(app_state, model_run_state)
elif model_run_state.chunking_strategy == CHUNKING_STRATEGY_SEMANTIC:
combined_document_objects = chunk_with_semantic_splitter(app_state, model_run_state)
else:
combined_document_objects = chunk_with_recursive_splitter(app_state, model_run_state)
model_run_state.combined_document_objects = combined_document_objects
dprint(app_state, "Chunking completed successfully")
def chunk_with_recursive_splitter(app_state, model_run_state):
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=model_run_state.chunk_size,
chunk_overlap=model_run_state.chunk_overlap,
length_function = tiktoken_len,
)
combined_document_objects = []
dprint(app_state, "Chunking documents and creating document objects")
for document in app_state.documents:
dprint(app_state, f"processing documend: {document['title']}")
text = document["single_text_document"]
dprint(app_state, text)
title = document["title"]
# document_id = document["document_id"]
chunks_document = text_splitter.split_text(text)
dprint(app_state, len(chunks_document))
for chunk_number, chunk in enumerate(chunks_document, start=1):
document_objects = Document(
page_content=chunk,
metadata={
"source": title,
"document_id": document.get("document_id", "default_id"),
"chunk_number": chunk_number # Add unique chunk number
}
)
combined_document_objects.append(document_objects)
return combined_document_objects
def chunk_with_table_aware(app_state, model_run_state):
combined_document_objects = []
dprint(app_state, "Using Table-Aware Chunking for documents.")
for document in app_state.documents:
title = document["title"]
text = document["single_text_document"]
# Check if document is a PDF and contains tables
if document.get("is_pdf", False):
with pdfplumber.open(document["file_path"]) as pdf:
for page in pdf.pages:
tables = page.extract_tables()
for table in tables:
table_content = "\n".join([str(row) for row in table])
document_objects = Document(
page_content=table_content,
metadata={
"source": title,
"document_id": document.get("document_id", "default_id"),
"chunk_number": "table"
}
)
combined_document_objects.append(document_objects)
# Chunk the rest of the text
text_splitter = RecursiveCharacterTextSplitter(chunk_size=model_run_state.chunk_size, chunk_overlap=model_run_state.chunk_overlap)
chunks_document = text_splitter.split_text(text)
for chunk_number, chunk in enumerate(chunks_document, start=1):
document_objects = Document(
page_content=chunk,
metadata={
"source": title,
"document_id": document.get("document_id", "default_id"),
"chunk_number": chunk_number
}
)
combined_document_objects.append(document_objects)
return combined_document_objects
def chunk_with_section_based(app_state, model_run_state):
combined_document_objects = []
dprint(app_state, "Using Section-Based Chunking for documents.")
for document in app_state.documents:
text = document["single_text_document"]
title = document["title"]
# Split the text by headings
sections = re.split(r"\n[A-Z].+?\n", text)
# Chunk each section
text_splitter = RecursiveCharacterTextSplitter(chunk_size=model_run_state.chunk_size, chunk_overlap=model_run_state.chunk_overlap)
for section_number, section in enumerate(sections, start=1):
chunks_document = text_splitter.split_text(section)
for chunk_number, chunk in enumerate(chunks_document, start=1):
document_objects = Document(
page_content=chunk,
metadata={
"source": title,
"document_id": document.get("document_id", "default_id"),
"section_number": section_number,
"chunk_number": chunk_number
}
)
combined_document_objects.append(document_objects)
return combined_document_objects
def chunk_with_semantic_splitter(app_state, model_run_state):
# Load pre-trained model for embeddings
model = SentenceTransformer('all-MiniLM-L6-v2')
combined_document_objects = []
dprint(app_state, "Using Semantic-Based Chunking for documents.")
for document in app_state.documents:
text = document["single_text_document"]
title = document["title"]
# Split text into sentences or paragraphs
sentences = text.split(". ") # Simple split by sentence (you can refine this)
sentence_embeddings = model.encode(sentences)
# Group sentences into chunks based on semantic similarity
chunks = []
current_chunk = []
for i in range(len(sentences) - 1):
current_chunk.append(sentences[i])
# Calculate similarity between consecutive sentences
sim = cosine_similarity([sentence_embeddings[i]], [sentence_embeddings[i + 1]])[0][0]
if sim < 0.7 or len(current_chunk) >= model_run_state.chunk_size:
# If similarity is below threshold or chunk size is reached, start a new chunk
chunks.append(" ".join(current_chunk))
current_chunk = []
# Add the final chunk
if current_chunk:
chunks.append(" ".join(current_chunk))
# Create document objects for the chunks
for chunk_number, chunk in enumerate(chunks, start=1):
document_objects = Document(
page_content=chunk,
metadata={
"source": title,
"document_id": document.get("document_id", "default_id"),
"chunk_number": chunk_number
}
)
combined_document_objects.append(document_objects)
return combined_document_objects |