add : select embedding_model

seohyunjun · Mar 10, 2024 · 5cbfa9a · 5cbfa9a
1 parent 45fa2af
commit 5cbfa9a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/main.py b/main.py
@@ -25,11 +25,11 @@ def main(config):
     args.add_argument("--chunk_size", type=int, default=5000)
     args.add_argument("--outputfile", type=str, default='test.md')
     args.add_argument("--model", type=str, default="gpt-3.5-turbo-16k")
-    args.add_argument("--title", type=str, default="title")
+    args.add_argument("--embedding_model", type=str, default="text-embedding-3-small")
+    args.add_argument("--title", type=str)
     # pinecone args
     args.add_argument("--pinecone", type=bool, default=False)
     args.add_argument("--chromadb", type=bool, default=True)
-    args.add_argument("--index_name", type=str, default="knowledge")
     config = args.parse_args()
 
     main(config)
diff --git a/utils/splitter.py b/utils/splitter.py
@@ -53,10 +53,10 @@ def splitter(config):
         # embedding document and upload to pinecone
 
         # emebdding document
-        pinecone_embed(index_name=config.title, docs=config.document)
+        pinecone_embed(index_name=config.title, model=config.embedding_model, docs=config.document)
     if config.chromadb:
         # embedding document and upload to pinecone
 
         # emebdding document
-        chromadb_embed(collection_name=config.title, docs=config.document)
+        chromadb_embed(collection_name=config.title, docs=config.document, model=config.embedding_model, embed_batch_size=100)