diff --git a/duas_query.py b/duas_query.py index 9c786e4..5956e3a 100644 --- a/duas_query.py +++ b/duas_query.py @@ -1,4 +1,4 @@ -from langchain_openai import OpenAIEmbeddings +from langchain_huggingface import HuggingFaceEmbeddings from dotenv import load_dotenv from langchain_postgres import PGVector import os @@ -9,7 +9,7 @@ COLLECTION_NAME = os.getenv('COLLECTION_NAME') # Initialize embeddings (needed for query encoding only) -embeddings = OpenAIEmbeddings(model="text-embedding-3-small") +embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2") # Load existing vector store db = PGVector( diff --git a/generate_dua_tags_embedding.py b/generate_dua_tags_embedding.py index 103f373..d95d37f 100644 --- a/generate_dua_tags_embedding.py +++ b/generate_dua_tags_embedding.py @@ -1,5 +1,5 @@ import json -from langchain_openai import OpenAIEmbeddings +from langchain_huggingface import HuggingFaceEmbeddings # from langchain.vectorstores.pgvector import PGVector from langchain_postgres import PGVector # from langchain.schema import Document @@ -56,7 +56,7 @@ for dua in duas_data: print(f"Created {len(documents)} documents from tags") # Initialize embeddings -embeddings = OpenAIEmbeddings(model="text-embedding-3-small") +embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2") # Create vector store - embeddings will be created only from tags (page_content) print("Creating embeddings and storing in pgvector...") diff --git a/main.py b/main.py index b8a4549..c44d89f 100644 --- a/main.py +++ b/main.py @@ -2,7 +2,7 @@ from fastapi import FastAPI, HTTPException, Query from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field from typing import List, Optional -from langchain_openai import OpenAIEmbeddings +from langchain_huggingface import HuggingFaceEmbeddings from langchain_postgres import PGVector from dotenv import load_dotenv import uvicorn @@ -32,7 +32,7 @@ COLLECTION_NAME = os.getenv('COLLECTION_NAME') # Initialize embeddings and vector store -embeddings = OpenAIEmbeddings(model="text-embedding-3-small") +embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2") db = PGVector( collection_name=COLLECTION_NAME, connection=CONNECTION_STRING,