Update app.py

app.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI
 import os
-import pymupdf  # PyMuPDF
+import fitz  # PyMuPDF
 from pptx import Presentation
 from sentence_transformers import SentenceTransformer
 import torch
@@ -11,18 +11,18 @@ import numpy as np
 
 app = FastAPI()
 
-# Initialize ChromaDB
+# Initialize ChromaDB with 512 dimensions
 client = chromadb.PersistentClient(path="/data/chroma_db")
-collection = client.get_or_create_collection(name="knowledge_base")
+collection = client.get_or_create_collection(name="knowledge_base", metadata={"hnsw:space": "cosine"}, embedding_function=None)
 
 # File Paths
 pdf_file = "Sutures and Suturing techniques.pdf"
 pptx_file = "impalnt 1.pptx"
 
 # Initialize Embedding Models
-text_model = SentenceTransformer('paraphrase-MiniLM-L12-v2')  # 512D embeddings
-model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+text_model = SentenceTransformer('paraphrase-MiniLM-L12-v2')  # 512D text embeddings
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
 # Image Storage Folder
 IMAGE_FOLDER = "/data/extracted_images"
@@ -31,7 +31,7 @@ os.makedirs(IMAGE_FOLDER, exist_ok=True)
 # Extract Text from PDF
 def extract_text_from_pdf(pdf_path):
     try:
-        doc = pymupdf.open(pdf_path)
+        doc = fitz.open(pdf_path)
         text = " ".join(page.get_text() for page in doc)
         return text.strip() if text else None
     except Exception as e:
@@ -53,7 +53,7 @@ def extract_text_from_pptx(pptx_path):
 # Extract Images from PDF
 def extract_images_from_pdf(pdf_path):
     try:
-        doc = pymupdf.open(pdf_path)
+        doc = fitz.open(pdf_path)
         images = []
         for i, page in enumerate(doc):
             for img_index, img in enumerate(page.get_images(full=True)):
@@ -93,9 +93,9 @@ def get_text_embedding(text):
 def get_image_embedding(image_path):
     try:
         image = Image.open(image_path)
-        inputs = processor(images=image, return_tensors="pt")
+        inputs = clip_processor(images=image, return_tensors="pt")
         with torch.no_grad():
-            image_embedding = model.get_image_features(**inputs).numpy().flatten()
+            image_embedding = clip_model.get_image_features(**inputs).squeeze().numpy()
         return image_embedding.tolist()
     except Exception as e:
         print(f"Error generating image embedding: {e}")