[feat] try to use embedding from vllm

Author: ThomasFaria

pyproject.toml

@@ -12,10 +12,10 @@ dependencies = [
     "langchain-neo4j>=0.4.0",
     "pandas>=2.2.3",
     "s3fs>=2024.12.0",
-    "langchain-huggingface>=0.1.2",
     "uvicorn>=0.34.0",
     "fastapi>=0.115.12",
     "streamlit>=1.44.0",
+    "langchain-openai>=0.3.11",
 ]
 authors = [
   {name="Thomas Faria", email="[email protected]"}

src/constants/graph_db.py

@@ -3,4 +3,5 @@
 NEO4J_URL = "neo4j://neo4j-585569.projet-ape:7687"
 NEO4J_USERNAME = "neo4j"
 NEO4J_PWD = os.environ["NEO4J_API_KEY"]
-EMBEDDING_MODEL = "OrdalieTech/Solon-embeddings-large-0.1"
+EMBEDDING_MODEL = "ordalieTech/Solon-embeddings-large-0.1"
+URL_EMBEDDING_API = "http://user-tfaria-vllm:8000/v1"

src/llm/client.py

@@ -8,7 +8,7 @@
 
 @asynccontextmanager
 async def get_llm_client():
-    client = AsyncOpenAI(api_key="EMPTY", base_url=URL_LLM_API, timeout=httpx.Timeout(30.0))
+    client = AsyncOpenAI(api_key="EMPTY", base_url=URL_LLM_API, timeout=httpx.Timeout(1 * 60 * 60))
     try:
         yield client
     finally:

src/vector_db/loaders.py

@@ -1,10 +1,10 @@
 import logging
 
-import torch
-from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_neo4j import Neo4jGraph, Neo4jVector
+from langchain_openai import OpenAIEmbeddings
 
-from constants.graph_db import EMBEDDING_MODEL, NEO4J_PWD, NEO4J_URL, NEO4J_USERNAME
+# from vector_db.openai_embeddings import CustomOpenAIEmbeddings
+from constants.graph_db import EMBEDDING_MODEL, NEO4J_PWD, NEO4J_URL, NEO4J_USERNAME, URL_EMBEDDING_API
 
 logger = logging.getLogger(__name__)
 
@@ -24,25 +24,16 @@ def setup_graph() -> Neo4jGraph:
     )
 
 
-def get_embedding_model(model_name: str) -> HuggingFaceEmbeddings:
-    """Initialize the HuggingFace embedding model."""
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    if device == "cpu":
-        logger.info("No GPU found: running on CPU. The embedding step might be slow 🫠")
-    elif device == "cuda":
-        logger.info("Running on GPU 🚀")
-
-    return HuggingFaceEmbeddings(
-        model_name=model_name,
-        model_kwargs={"device": device},
-        encode_kwargs={"normalize_embeddings": True},
-        show_progress=False,
+def get_embedding_model(model_name: str) -> OpenAIEmbeddings:
+    """Initialize the embedding model."""
+    return OpenAIEmbeddings(
+        open=model_name,
+        openai_base_url=URL_EMBEDDING_API,
+        openai_api_key="EMPTY",
     )
 
 
-def get_vector_db() -> Neo4jVector:
+async def get_vector_db() -> Neo4jVector:
     """Initialize the Neo4jVector Store from existing graph."""
     emb_model = get_embedding_model(EMBEDDING_MODEL)
     graph = setup_graph()