wassname/stampy-chat
1
0
Fork 0
mirror of https://github.com/wassname/stampy-chat.git synced 2026-09-11 12:50:34 +08:00
Code Issues Packages Projects Releases Wiki Activity

Added sentence transformer option

Browse Source
This commit is contained in:
henri123lemoine
2023-07-02 07:09:06 -04:00
parent 0cf2fdad74
commit cdfb5a52a7
2 changed files with 20 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/dataset/settings.py
+4 -2
View File
@@ -11,9 +11,11 @@ SQL_DB_PATH = str(current_file_path.parent / 'data' / 'ARD.db')
ARD_DATASET_NAME = "StampyAI/alignment-research-dataset"
### EMBEDDINGS ###
EMBEDDINGS_MODEL = "text-embedding-ada-002"
USE_OPENAI_EMBEDDINGS = False
OPENAI_EMBEDDINGS_MODEL = "text-embedding-ada-002"
SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL = "sentence-transformers/multi-qa-mpnet-base-cos-v1"
EMBEDDINGS_DIMS = 1536
EMBEDDINGS_RATE_LIMIT = 3500
OPENAI_EMBEDDINGS_RATE_LIMIT = 3500
### PINECONE ###
PINECONE_INDEX_NAME = "stampy-chat-embeddings-test"
src/dataset/update_dataset.py
+16 -7
View File
@@ -11,7 +11,7 @@ from .text_splitter import TokenSplitter
from .sql_db_handler import SQLDB
from .pinecone_db_handler import PineconeDB
from .settings import EMBEDDINGS_MODEL, EMBEDDINGS_DIMS, EMBEDDINGS_RATE_LIMIT, ARD_DATASET_NAME, MAX_NUM_AUTHORS_IN_SIGNATURE
from .settings import USE_OPENAI_EMBEDDINGS, OPENAI_EMBEDDINGS_MODEL, SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL, EMBEDDINGS_DIMS, OPENAI_EMBEDDINGS_RATE_LIMIT, ARD_DATASET_NAME, MAX_NUM_AUTHORS_IN_SIGNATURE
import logging
logger = logging.getLogger(__name__)
@@ -26,6 +26,12 @@ class ARDUpdater:
self.token_splitter = TokenSplitter(min_tokens_per_block, max_tokens_per_block)
self.sql_db = SQLDB()
self.pinecone_db = PineconeDB()
if not USE_OPENAI_EMBEDDINGS:
from langchain.embeddings import HuggingFaceEmbeddings
self.hf_embeddings = HuggingFaceEmbeddings(
model_name=SENTENCE_TRANSFORMER_EMBEDDINGS_MODEL,
)
def update(self, custom_sources: List[str] = ['all']):
for source in custom_sources:
@@ -48,11 +54,14 @@ class ARDUpdater:
chunks_ids_batch = batch['chunks_ids_batch']
try:
embeddings = self.get_embeddings(chunks_batch)
if USE_OPENAI_EMBEDDINGS:
embeddings = self.get_openai_embeddings(chunks_batch)
else:
embeddings = np.array(self.hf_embeddings.embed_documents(chunks_batch))
self.sql_db.upsert_chunks(chunks_ids_batch, chunks_batch)
self.pinecone_db.delete_entries([entry['id'] for entry in entries_batch])
self.pinecone_db.upsert_entries(entries_batch, chunks_batch, chunks_ids_batch, embeddings)
# self.pinecone_db.delete_entries([entry['id'] for entry in entries_batch])
# self.pinecone_db.upsert_entries(entries_batch, chunks_batch, chunks_ids_batch, embeddings)
logger.info(f"Successfully updated {len(entries_batch)} {source} entries with {len(chunks_batch)} total chunks.")
except Exception as e:
@@ -132,12 +141,12 @@ class ARDUpdater:
return self.sql_db.upsert_entry(entry)
@retry(stop=stop_after_attempt(3))
def get_embeddings(self, chunks):
def get_openai_embeddings(self, chunks):
embeddings = np.zeros((len(chunks), EMBEDDINGS_DIMS))
rate_limit = EMBEDDINGS_RATE_LIMIT # TODO: use this rate_limit
rate_limit = OPENAI_EMBEDDINGS_RATE_LIMIT # TODO: use this rate_limit
openai_output = openai.Embedding.create(
model=EMBEDDINGS_MODEL,
model=OPENAI_EMBEDDINGS_MODEL,
input=chunks
)['data']