fastapiapp / file_processing.py
Sk4467's picture
added application
d0fbfa7
raw
history blame
3.75 kB
from langchain.document_loaders import PyPDFLoader, PDFMinerLoader, DirectoryLoader
from langchain.embeddings import SentenceTransformerEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from os.path import join
import os
from dotenv import load_dotenv
load_dotenv(r'C:\Users\sksha\Desktop\llm-assignment-master\llm-assignment-master\llm-assignment-master_\backend\.env')
openai_api_key = os.environ.get('OPENAI_API_KEY')
from langchain.document_loaders import TextLoader, PDFMinerLoader, UnstructuredWordDocumentLoader, CSVLoader
# def load_documents(file_path):
# if file_path.endswith('.txt'):
# loader = TextLoader(file_path)
# elif file_path.endswith('.pdf'):
# loader = PyPDFLoader(file_path)
# elif file_path.endswith('.doc') or file_path.endswith('.docx'):
# loader = UnstructuredWordDocumentLoader(file_path)
# elif file_path.endswith('.csv'):
# loader = CSVLoader(file_path)
# else:
# raise ValueError(f"Unsupported file format: {file_path}")
# documents = loader.load()
# return documents
from fastapi import UploadFile
from typing import List
import fitz # PyMuPDF
import pandas as pd
import docx
from langchain.docstore.document import Document
def read_pdf(file_path: str) -> str:
doc = fitz.open(file_path)
text = ""
for page in doc:
text += page.get_text()
return text
def read_docx(file_path: str) -> str:
doc = docx.Document(file_path)
fullText = []
for para in doc.paragraphs:
fullText.append(para.text)
return '\n'.join(fullText)
def read_csv(file_path: str) -> str:
df = pd.read_csv(file_path)
return df.to_string()
def read_txt(file_path: str) -> str:
with open(file_path, 'r', encoding='utf-8') as file:
return file.read()
async def load_documents(file: UploadFile)->List[Document]:
temp_file_path = f"temp_{file.filename}"
try:
# Save the uploaded file to a temporary file
with open(temp_file_path, "wb") as temp_file:
temp_file.write(await file.read())
content = ""
if file.filename.endswith('.pdf'):
content = read_pdf(temp_file_path)
elif file.filename.endswith('.docx'):
content = read_docx(temp_file_path)
elif file.filename.endswith('.csv'):
content = read_csv(temp_file_path)
elif file.filename.endswith('.txt'):
content = read_txt(temp_file_path)
else:
raise ValueError("Unsupported file format")
except Exception as e:
# Handle general errors - log or adjust as necessary for your application
print(f"Error processing document: {e}")
content = "Error processing document."
finally:
# Cleanup: remove the temporary file
if os.path.exists(temp_file_path):
os.remove(temp_file_path)
metadata = {'source': file.filename}
document = Document(page_content=content, metadata=metadata)
return [document]
from langchain.text_splitter import CharacterTextSplitter
def chunk_documents(documents, chunk_size=1000, chunk_overlap=200):
text_splitter = CharacterTextSplitter(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap
)
chunked_docs = text_splitter.split_documents(documents)
return chunked_docs
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
def create_embeddings(chunked_docs, collection_name):
embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key)
vector_store = Chroma.from_documents(chunked_docs, embeddings, collection_name=collection_name)
vector_store.persist()
return vector_store