Team Ai
Apppublic

Multimedika/Bot_Development

sourceHugging Facemitupdated 2y agoView on Hugging Face
0likes
document_uploader.py117 linesDownload Raw Back to script
1from llama_index.core.ingestion import IngestionPipeline2from llama_index.embeddings.openai import OpenAIEmbedding3from config import PINECONE_CONFIG4from pinecone.grpc import PineconeGRPC as Pinecone5# from service.reader import Reader6from script.get_metadata import Metadata7from fastapi import UploadFile, status8from fastapi.responses import JSONResponse9 10from llama_index.core.node_parser import (11    SentenceSplitter,12    SemanticSplitterNodeParser,13)14from llama_index.core import Settings15# from service.reader_v3 import upload_file16from service.reader_v4 import upload_file17 18# from script.get_topic import extract_topic19 20import logging21import random22 23 24class Uploader:25    # def __init__(self, reference, file: UploadFile, content_table: UploadFile):26    def __init__(self, reference, file: UploadFile, lang: str = "en"):27        self.file = file28        # self.content_table = content_table29        # self.reader = Reader()30        self.reference = reference31        self.metadata = Metadata(reference)32        self.lang = lang33 34    def check_existing_metadata(self, pinecone_index, title, random_vector):35        try:36            result = pinecone_index.query(37                vector=random_vector,38                top_k=1,39                filter={40                    "title": {"$eq": title},41                },42            )43            return result["matches"]44        except Exception as e:45            return JSONResponse(46                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,47                content=f"Error check existing metadata {str(e)}",48            )49 50    async def process_documents(self):51        # Get metadata52        documents_with_metadata, file_stream = await upload_file(self.reference, self.file, self.lang)53 54        if isinstance(documents_with_metadata, JSONResponse):55            return documents_with_metadata  # Return the error response directly56 57 58        # embed_model = OpenAIEmbedding()59        embed_model = OpenAIEmbedding(model="text-embedding-3-large")60        Settings.embed_model = embed_model61        # Set up the ingestion pipeline62        pipeline = IngestionPipeline(63            transformations=[64                SemanticSplitterNodeParser(65                    buffer_size=1,66                    breakpoint_percentile_threshold=95,67                    embed_model=embed_model,68                ),69                # topic_extractor,70            ]71        )72 73        # splitter = SemanticSplitterNodeParser(74        #     buffer_size=1, breakpoint_percentile_threshold=95, embed_model=embed_model75        # )76 77        #  Run the pipeline78        try:79            print("Pipeline processing completed with Semantic Spliter.")80            nodes_with_metadata = pipeline.run(documents=documents_with_metadata)81            return nodes_with_metadata, file_stream82 83        except Exception as e:84            try:85                # If the first method fails, fallback to sentence splitter86                sentence_splitter = SentenceSplitter(chunk_size=512)87                nodes_with_metadata = sentence_splitter.get_nodes_from_documents(88                    documents_with_metadata89                )90                print("Pipeline processing completed with SentenceSplitter fallback.")91                return nodes_with_metadata, file_stream92            except Exception as fallback_error:93                # Log the second error and return JSONResponse for FastAPI94                logging.error(f"Error with SentenceSplitter fallback: {fallback_error}")95                return JSONResponse(96                    status_code=500,97                    content="An internal server error occurred during pipeline processing.",98                )99 100    def filter_document(self, documents):101        api_key = PINECONE_CONFIG.PINECONE_API_KEY102        client = Pinecone(api_key=api_key)103        pinecone_index = client.Index("test")104 105        random_vector = [random.uniform(0, 1) for _ in range(1536)]106 107        filtered_documents = []108        for doc in documents:109            result = self.check_existing_metadata(110                pinecone_index, doc.metadata["title"], random_vector111            )112 113            if len(result) == 0:114                filtered_documents.append(doc)115 116        return filtered_documents117