Update vector_store_retriever.py

vector_store_retriever.py

@@ -3,32 +3,51 @@ import os
 import gradio as gr
 import time
 from pydantic import BaseModel, Field
-from typing import Any, Optional, Dict, List
+from typing import Any, Optional, Dict, List, Union
 from huggingface_hub import InferenceClient
-from langchain.llms.base import LLM
-from langchain.embeddings import HuggingFaceInstructEmbeddings
+from langchain.llms.base import LLM, Documents, Images, EmbeddingFunction, Embeddings
 from langchain.vectorstores import Chroma
 from dotenv import load_dotenv
-from transformers import AutoTokenizer
-from transformers import Tool
+from transformers import AutoTokenizer, AutoModel, Tool
 
 load_dotenv()
 
 path_work = "."
 hf_token = os.getenv("HF")
 
-embeddings = HuggingFaceInstructEmbeddings(
-    model_name="sentence-transformers/all-MiniLM-L6-v2",
-    model_kwargs={"device": "cpu"}
-)
+class HuggingFaceInstructEmbeddings(EmbeddingFunction):
+    def __init__(self, model_name: str, model_kwargs: Optional[Dict[str, Any]] = None):
+        self.model = AutoModel.from_pretrained(model_name, **(model_kwargs or {}))
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    def __call__(self, input: Union[Documents, Images]) -> Embeddings:
+        if isinstance(input, Documents):
+            texts = [doc.text for doc in input]
+            embeddings = self._embed_text(texts)
+        else:
+            # Handle image embeddings if needed
+            pass
+
+        return embeddings
+
+    def _embed_text(self, texts: List[str]) -> Embeddings:
+        # Your existing logic for text embeddings using Hugging Face models...
+        inputs = self.tokenizer(texts, return_tensors="pt", padding=True, truncation=True)
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+        embeddings = outputs.last_hidden_state.mean(dim=1)  # Adjust this based on your specific model
+
+        return embeddings
+
 
 vectordb = Chroma(
     persist_directory=path_work + '/new_papers',
-    embedding_function=embeddings
+    embedding_function=HuggingFaceInstructEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={"device": "cpu"})
 )
 
 retriever = vectordb.as_retriever(search_kwargs={"k": 2})#5
 
+
 class KwArgsModel(BaseModel):
     kwargs: Dict[str, Any] = Field(default_factory=dict)