From 3fb716b2a621eca8dab8ad4671c7c651d4043787 Mon Sep 17 00:00:00 2001
From: moon <shawmakesmagic@gmail.com>
Date: Wed, 9 Aug 2023 23:55:45 -0700
Subject: [PATCH] change sentence transformer model

---
 agentmemory/postgres.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/agentmemory/postgres.py b/agentmemory/postgres.py
index 56e5a54..d218999 100644
--- a/agentmemory/postgres.py
+++ b/agentmemory/postgres.py
@@ -1,6 +1,5 @@
 import json
 import psycopg2
-from sentence_transformers import SentenceTransformer
 
 
 class PostgresCollection:
@@ -157,7 +156,10 @@ def __init__(self, connection_string):
         from pgvector.psycopg2 import register_vector
 
         register_vector(self.cur)  # Register PGVector functions
-        self.model = SentenceTransformer("BAAI/bge-large-en")
+
+        from sentence_transformers import SentenceTransformer
+
+        self.model = SentenceTransformer("all-MiniLM-L6-v2")
 
     def _table_name(self, category):
         return f"memory_{category}"
@@ -264,7 +266,7 @@ def query(self, category, query_texts, n_results=5):
             "documents": [],
             "metadatas": [],
             "embeddings": [],
-            "distances": []
+            "distances": [],
         }
         self.ensure_table_exists(category)
         table_name = self._table_name(category)
@@ -288,7 +290,6 @@ def query(self, category, query_texts, n_results=5):
                     results["distances"].append(row[4])
         return results
 
-
     def update(self, category, id_, document=None, metadata=None, embedding=None):
         self.ensure_table_exists(category)
         table_name = self._table_name(category)