fix issues caused by refactor

CambioML · Sep 13, 2023 · 2331aa5 · 2331aa5
1 parent 5b183e8
commit 2331aa5
Show file tree

Hide file tree

Showing 6 changed files with 12 additions and 3,799 deletions.
diff --git a/example/flex/flex_demo.py b/example/flex/flex_demo.py
@@ -2,13 +2,17 @@
 
 import os
 import argparse
+from dotenv import load_dotenv
 from pykoi import Application
 from pykoi.chat import QuestionAnswerDatabase
 from pykoi.retrieval import RetrievalFactory
 from pykoi.retrieval import VectorDbFactory
 from pykoi.component import Chatbot, Dashboard, RetrievalQA
 
 
+load_dotenv()
+
+
 def main(**kargs):
     os.environ["DOC_PATH"] = os.path.join(os.getcwd(), "temp/docs")
     os.environ["VECTORDB_PATH"] = os.path.join(os.getcwd(), "temp/vectordb")

diff --git a/poetry.lock b/poetry.lock
diff --git a/pykoi/retrieval/llm/embedding_factory.py b/pykoi/retrieval/llm/embedding_factory.py
@@ -27,8 +27,10 @@ def create_embedding(model_source: Union[str, ModelSource], **kwargs) -> Embeddi
         try:
             model_source = ModelSource(model_source)
             if model_source == ModelSource.OPENAI:
+                from langchain.embeddings import OpenAIEmbeddings
                 return OpenAIEmbeddings()
             elif model_source == ModelSource.HUGGINGFACE:
+                from langchain.embeddings import HuggingFaceEmbeddings
                 return HuggingFaceEmbeddings(
                     model_name=kwargs.get("model_name"),
                 )

diff --git a/pykoi/retrieval/llm/openai.py b/pykoi/retrieval/llm/openai.py
@@ -5,12 +5,10 @@
 
 from langchain.chains import RetrievalQA
 from langchain.llms import OpenAI
-from dotenv import load_dotenv
 
 from pykoi.retrieval.llm.abs_llm import AbsLlm
 from pykoi.retrieval.vectordb.abs_vectordb import AbsVectorDb
 
-load_dotenv()
 
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 MIN_DOCS = 2

diff --git a/pykoi/retrieval/llm/retrieval_factory.py b/pykoi/retrieval/llm/retrieval_factory.py
@@ -2,8 +2,6 @@
 from typing import Union
 
 from pykoi.retrieval.llm.abs_llm import AbsLlm
-from pykoi.retrieval.llm.openai import OpenAIModel
-from pykoi.retrieval.llm.huggingface import HuggingFaceModel
 from pykoi.retrieval.llm.constants import ModelSource
 from pykoi.retrieval.vectordb.abs_vectordb import AbsVectorDb
 
@@ -29,8 +27,10 @@ def create(
         try:
             model_source = ModelSource(model_source)
             if model_source == ModelSource.OPENAI:
+                from pykoi.retrieval.llm.openai import OpenAIModel
                 return OpenAIModel(vector_db)
             if model_source == ModelSource.HUGGINGFACE:
+                from pykoi.retrieval.llm.huggingface import HuggingFaceModel
                 return HuggingFaceModel(vector_db, **kwargs)
         except Exception as ex:
             raise Exception(f"Unknown model: {model_source}") from ex
diff --git a/pyproject.toml b/pyproject.toml
@@ -21,6 +21,7 @@ bcrypt = "4.0.1"
 posthog = "3.0.1"
 pynvml = "11.5.0"
 pandas = "2.0.3"
+python-dotenv = "1.0.0"
 
 transformers = { version = "4.31.0", optional = true }
 einops = { version = "0.6.1", optional = true }
@@ -47,7 +48,8 @@ huggingface = [
     "transformers",
     "einops",
     "accelerate",
-    "bitsandbytes"
+    "bitsandbytes",
+    "sentence-transformers"
 ]
 rag = [
     "langchain", 
@@ -57,8 +59,7 @@ rag = [
     "pdfminer-six", 
     "docx2txt", 
     "python-multipart", 
-    "tiktoken", 
-    "sentence-transformers"
+    "tiktoken"
 ]
 rlhf = [
     "transformers",