Upload 2 files

.gitattributes

@@ -38,3 +38,4 @@ vol2.pdf filter=lfs diff=lfs merge=lfs -text
 vol3.pdf filter=lfs diff=lfs merge=lfs -text
 vol4.pdf filter=lfs diff=lfs merge=lfs -text
 vol5.pdf filter=lfs diff=lfs merge=lfs -text
+test.pdf filter=lfs diff=lfs merge=lfs -text

test.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ecd8e1207b3be0e246d40823509a2c774594319601bd1c28171722f735058a2e
+size 381011

train.py

@@ -0,0 +1,137 @@
+import fitz  # PyMuPDF
+from transformers import AutoTokenizer, AutoModel
+import torch
+from milvus import Milvus, DataType
+import os
+from langchain_community.llms import HuggingFaceEndpoint
+from langchain.prompts import PromptTemplate
+from langchain.schema import StrOutputParser
+from langchain.schema.runnable import RunnablePassthrough
+from langchain.schema import BaseLoader, LCDocument
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from docling.document_converter import DocumentConverter
+import gradio as gr
+from typing import Iterator
+
+# Initialize Milvus
+milvus = Milvus(host='localhost', port='19530')
+
+# Load BAAI embedding model
+tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-base-en-v1.5")
+model = AutoModel.from_pretrained("BAAI/bge-base-en-v1.5")
+
+# Docling PDF Loader
+class DoclingPDFLoader(BaseLoader):
+    def __init__(self, file_path: str | list[str]) -> None:
+        self._file_paths = file_path if isinstance(file_path, list) else [file_path]
+        self._converter = DocumentConverter()
+
+    def lazy_load(self) -> Iterator[LCDocument]:
+        for source in self._file_paths:
+            dl_doc = self._converter.convert(source).document
+            text = dl_doc.export_to_markdown()
+            yield LCDocument(page_content=text)
+
+    def load(self) -> list[LCDocument]:
+        return list(self.lazy_load())
+
+# Function to extract and split text from PDF
+def extract_text_from_pdf(pdf_path):
+    loader = DoclingPDFLoader(file_path=pdf_path)
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1000,
+        chunk_overlap=200,
+    )
+    docs = loader.load()
+    splits = text_splitter.split_documents(docs)
+    return " ".join([doc.page_content for doc in splits])
+
+# Set up LLM
+HF_API_KEY = os.environ.get("HF_API_KEY")
+HF_LLM_MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.3"
+llm = HuggingFaceEndpoint(
+    repo_id=HF_LLM_MODEL_ID,
+    huggingfacehub_api_token=HF_API_KEY,
+)
+
+# Function to extract text from PDF
+def extract_text_from_pdf(pdf_path):
+    doc = fitz.open(pdf_path)
+    text = ""
+    for page in doc:
+        text += page.get_text()
+    return text
+
+# Function to generate embeddings
+def generate_embeddings(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    return outputs.last_hidden_state.mean(dim=1).squeeze().numpy()
+
+# Function to insert embeddings into Milvus
+def insert_into_milvus(embeddings):
+    collection_name = "pdf_embeddings"
+    if not milvus.has_collection(collection_name):
+        milvus.create_collection({
+            "collection_name": collection_name,
+            "dimension": embeddings.shape[0],
+            "index_file_size": 1024,
+            "metric_type": "L2"
+        })
+    milvus.insert(collection_name, [embeddings])
+
+# Function to query Milvus
+def query_milvus(query_embedding, top_k=5):
+    collection_name = "pdf_embeddings"
+    search_params = {"metric_type": "L2", "params": {"nprobe": 16}}
+    results = milvus.search(collection_name, [query_embedding], top_k, search_params)
+    return results
+
+# Function to generate response using Llama
+# Update generate_response function to use the RAG pipeline
+def generate_response(query, context):
+    prompt = PromptTemplate.from_template(
+        "Context information is below.\n---------------------\n{context}\n---------------------\nGiven the context information and not prior knowledge, answer the query.\nQuery: {question}\nAnswer:\n"
+    )
+    
+    rag_chain = (
+        {"context": context, "question": query}
+        | prompt
+        | llm
+        | StrOutputParser()
+    )
+    
+    return rag_chain.invoke(query)
+
+# Main function
+def main(pdf_path, query):
+    # Step 1: Extract text from PDF
+    text = extract_text_from_pdf(pdf_path)
+    
+    # Step 2: Generate embeddings for the text
+    embeddings = generate_embeddings(text)
+    
+    # Step 3: Insert embeddings into Milvus
+    insert_into_milvus(embeddings)
+    
+    # Step 4: Generate embeddings for the query
+    query_embedding = generate_embeddings(query)
+    
+    # Step 5: Query Milvus for similar embeddings
+    results = query_milvus(query_embedding)
+    
+    # Step 6: Generate response using Llama
+    context = " ".join([result for result in results])
+    response = generate_response(query, context)
+    
+    print(response)
+
+if __name__ == "__main__":
+    # Initialize Gradio interface
+    def ask_question(question):
+        pdf_path = "test.pdf"
+        main(pdf_path, question)
+    
+    iface = gr.Interface(fn=ask_question, inputs="text", outputs="text")
+    iface.launch()