updated for torch

app/utils/vector_db.py

@@ -1,9 +1,12 @@
 import os
 import faiss
 import numpy as np
+import torch
 from sentence_transformers import SentenceTransformer
 from app.utils.file_handler import extract_text_from_file
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+#obj_embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2', device=device)
 # Load a pre-trained embedding model with GPU support
 obj_embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2', device="cuda")
 

requirements.txt

@@ -1,9 +1,9 @@
 Flask==2.3.2
 transformers==4.27.4
 sentence-transformers==2.2.2
-torch==2.0.1+cu118         # GPU-enabled PyTorch (CUDA 11.8 compatible)
-torchvision==0.15.2+cu118  # GPU-enabled TorchVision
-faiss-gpu==1.7.4           # GPU-enabled FAISS
+torch==2.0.1
+torchvision==0.15.2
+faiss-gpu==1.7.4
 PyMuPDF==1.22.5
 python-pptx==0.6.21
 pandas==1.5.3