File size: 5,056 Bytes
6d55408 |
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 |
import os
import hashlib
import io
import json
import pandas as pd
from langchain_community.vectorstores import FAISS
from PyPDF2 import PdfReader
from docx import Document
from langchain_huggingface import HuggingFaceEmbeddings
class FileHandler:
def __init__(self, vector_db_path,api_token):
self.vector_db_path = vector_db_path
# Initialize the embedding model using Hugging Face
self.embeddings = HuggingFaceEmbeddings(
model_name="sentence-transformers/all-MiniLM-L6-v2",
model_kwargs={"token": api_token},
)
def handle_file_upload(self, file, document_name, document_description):
try:
content = file.read()
file_hash = hashlib.md5(content).hexdigest()
file_key = f"{file.name}_{file_hash}"
vector_store_dir = os.path.join(self.vector_db_path, file_key)
os.makedirs(vector_store_dir, exist_ok=True)
vector_store_path = os.path.join(vector_store_dir, "index.faiss")
if os.path.exists(vector_store_path):
return {"message": "File already processed."}
# Process file based on type
if file.name.endswith(".pdf"):
texts, metadatas = self.load_and_split_pdf(file)
elif file.name.endswith(".docx"):
texts, metadatas = self.load_and_split_docx(file)
elif file.name.endswith(".txt"):
texts, metadatas = self.load_and_split_txt(content)
elif file.name.endswith(".xlsx"):
texts, metadatas = self.load_and_split_table(content)
elif file.name.endswith(".csv"):
texts, metadatas = self.load_and_split_csv(content)
else:
raise ValueError("Unsupported file format.")
if not texts:
return {"message": "No text extracted from the file. Check the file content."}
# Create FAISS vector store using LangChain's from_texts method
vector_store = FAISS.from_texts(texts, embedding=self.embeddings, metadatas=metadatas)
vector_store.save_local(vector_store_dir)
metadata = {
"filename": file.name,
"document_name": document_name,
"document_description": document_description,
"file_size": len(content),
}
metadata_path = os.path.join(vector_store_dir, "metadata.json")
with open(metadata_path, 'w') as md_file:
json.dump(metadata, md_file)
return {"message": "File processed successfully."}
except Exception as e:
return {"message": f"Error processing file: {str(e)}"}
def load_and_split_pdf(self, file):
reader = PdfReader(file)
texts = []
metadatas = []
for page_num, page in enumerate(reader.pages):
text = page.extract_text()
if text:
texts.append(text)
metadatas.append({"page_number": page_num + 1})
return texts, metadatas
def load_and_split_docx(self, file):
doc = Document(file)
texts = []
metadatas = []
for para_num, paragraph in enumerate(doc.paragraphs):
if paragraph.text:
texts.append(paragraph.text)
metadatas.append({"paragraph_number": para_num + 1})
return texts, metadatas
def load_and_split_txt(self, content):
text = content.decode("utf-8")
lines = text.split('\n')
texts = [line for line in lines if line.strip()]
metadatas = [{}] * len(texts)
return texts, metadatas
def load_and_split_table(self, content):
excel_data = pd.read_excel(io.BytesIO(content), sheet_name=None)
texts = []
metadatas = []
for sheet_name, df in excel_data.items():
df = df.dropna(how='all', axis=0).dropna(how='all', axis=1)
df = df.fillna('N/A')
for _, row in df.iterrows():
row_dict = row.to_dict()
# Combine key-value pairs into a string
row_text = ', '.join([f"{key}: {value}" for key, value in row_dict.items()])
texts.append(row_text)
metadatas.append({"sheet_name": sheet_name})
return texts, metadatas
def load_and_split_csv(self, content):
csv_data = pd.read_csv(io.StringIO(content.decode('utf-8')))
texts = []
metadatas = []
csv_data = csv_data.dropna(how='all', axis=0).dropna(how='all', axis=1)
csv_data = csv_data.fillna('N/A')
for _, row in csv_data.iterrows():
row_dict = row.to_dict()
row_text = ', '.join([f"{key}: {value}" for key, value in row_dict.items()])
texts.append(row_text)
metadatas.append({"row_index": _})
return texts, metadatas
|