parameterize embedding model and remove streamlit library

bilge-ince · Jul 31, 2024 · da3385b · da3385b
1 parent 8ae362e
commit da3385b
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 47 deletions.
diff --git a/.env-example b/.env-example
@@ -6,6 +6,7 @@ DB_HOST=localhost
 DB_PORT=5432
 
 # MODEl
+AIDB_MODEL_NAME=all-MiniLM-L6-v2
 MODEL_NAME=mistralai/Mistral-7B-Instruct-v0.2
 TOKENIZER_NAME=mistralai/Mistral-7B-Instruct-v0.2
 HUGGING_FACE_ACCESS_TOKEN=
diff --git a/embedding.py b/embedding.py
@@ -1,17 +1,17 @@
-# importing all the required modules
+import os
 import PyPDF2
 from db import get_connection
 
 def generate_embeddings():
     conn = get_connection()
     cursor = conn.cursor()
 
-    cursor.execute("""
+    cursor.execute(f"""
                     SELECT aidb.create_pg_retriever(
                         'documents_embeddings',
                         'public',
                         'id',
-                        'all-MiniLM-L6-v2',
+                        '{os.getenv("EMBEDDING_MODEL")}',
                         'text',
                         'documents',
                         ARRAY['id', 'doc_fragment'],

diff --git a/streamlit/chatgptui.py b/streamlit/chatgptui.py
diff --git a/streamlit/llamaindex.py b/streamlit/llamaindex.py