reformat well

.idea/llama-xpipe.iml

@@ -2,9 +2,10 @@
 <module type="PYTHON_MODULE" version="4">
   <component name="NewModuleRootManager">
     <content url="file://$MODULE_DIR$">
+      <excludeFolder url="file://$MODULE_DIR$/.venv" />
       <excludeFolder url="file://$MODULE_DIR$/venv" />
     </content>
-    <orderEntry type="jdk" jdkName="Python 3.10 (llama-xpipe)" jdkType="Python SDK" />
+    <orderEntry type="inheritedJdk" />
     <orderEntry type="sourceFolder" forTests="false" />
   </component>
 </module>

app.py

@@ -2,11 +2,12 @@ import logging
 import sys
 
 import streamlit as st
+from dotenv import load_dotenv
 
 from xpipe_wiki.manager_factory import XPipeRobotManagerFactory, XPipeRobotRevision
 
 logging.basicConfig(
-    stream=sys.stdout, level=logging.DEBUG
+    stream=sys.stdout, level=logging.INFO
 )  # logging.DEBUG for more verbose output
 logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
 
@@ -29,7 +30,7 @@ with st.sidebar:
 def main() -> None:
     st.header("X-Pipe Wiki 机器人 💬")
     robot_manager = XPipeRobotManagerFactory.get_or_create(
-        XPipeRobotRevision.SIMPLE_OPENAI_VERSION_0
+        XPipeRobotRevision.HUGGINGFACE_VERSION_0
     )
     robot = robot_manager.get_robot()
     query = st.text_input("X-Pipe Wiki 问题:")
@@ -39,4 +40,5 @@ def main() -> None:
 
 
 if __name__ == "__main__":
+    load_dotenv()
     main()

dataset/graph_store.json

@@ -1,3 +1 @@
-{
-  "graph_dict": {}
-}
+{"graph_dict": {}}

dataset/index_store.json

@@ -1,8 +1 @@
-{
-  "index_store/data": {
-    "7fd1f2d1-8a3a-4e24-bf42-5c149a73c8dc": {
-      "__type__": "vector_store",
-      "__data__": "{\"index_id\": \"7fd1f2d1-8a3a-4e24-bf42-5c149a73c8dc\", \"summary\": null, \"nodes_dict\": {\"190d619c-e92d-4cc5-b43b-652f66f4d6e0\": \"190d619c-e92d-4cc5-b43b-652f66f4d6e0\", \"1158d24f-29db-457b-aeaa-2a051316796e\": \"1158d24f-29db-457b-aeaa-2a051316796e\", \"520004a6-3eb6-44cf-90ba-a51ed1c6b894\": \"520004a6-3eb6-44cf-90ba-a51ed1c6b894\", \"359cdd4d-6096-4be4-991d-f3cc691c07c9\": \"359cdd4d-6096-4be4-991d-f3cc691c07c9\", \"3e4ffb91-090e-4e84-92cd-e8df8b3a87b8\": \"3e4ffb91-090e-4e84-92cd-e8df8b3a87b8\", \"7aa9b872-c4df-4de6-88c5-55962b913ba9\": \"7aa9b872-c4df-4de6-88c5-55962b913ba9\", \"16aa0efc-39b4-4fa8-93ab-f31b5740670c\": \"16aa0efc-39b4-4fa8-93ab-f31b5740670c\", \"1bebdce9-981b-4100-914d-48da73dc5f96\": \"1bebdce9-981b-4100-914d-48da73dc5f96\", \"bbe72c0d-ab9f-45c0-b581-d051f7aa797e\": \"bbe72c0d-ab9f-45c0-b581-d051f7aa797e\", \"229d36cc-1f60-485a-b44f-c53701b423de\": \"229d36cc-1f60-485a-b44f-c53701b423de\", \"dff86a2f-8cf4-4f01-a8ab-36977bb2db5c\": \"dff86a2f-8cf4-4f01-a8ab-36977bb2db5c\", \"cdf5fa07-4780-4cc2-9073-851996c4f006\": \"cdf5fa07-4780-4cc2-9073-851996c4f006\", \"a27c62d4-e969-46a8-94d8-aa0ce552234e\": \"a27c62d4-e969-46a8-94d8-aa0ce552234e\", \"e85816a0-a950-44c7-af93-dcf15a889490\": \"e85816a0-a950-44c7-af93-dcf15a889490\", \"cce4f38a-e5d4-4561-aab9-d695eef3bacf\": \"cce4f38a-e5d4-4561-aab9-d695eef3bacf\", \"b3716ae7-21ec-4112-979e-d5963c8235d5\": \"b3716ae7-21ec-4112-979e-d5963c8235d5\", \"55362a6d-27e6-46de-8fa8-c87a2039020f\": \"55362a6d-27e6-46de-8fa8-c87a2039020f\", \"785291f8-266d-4458-ba2c-2c110fac1038\": \"785291f8-266d-4458-ba2c-2c110fac1038\", \"c8655cf5-12f3-479a-a3b4-b43a2f6965ba\": \"c8655cf5-12f3-479a-a3b4-b43a2f6965ba\", \"67b2a33e-b9a3-4072-bc89-c1405d60dbe0\": \"67b2a33e-b9a3-4072-bc89-c1405d60dbe0\", \"42256b49-3c47-4fa2-8c64-871a757b9993\": \"42256b49-3c47-4fa2-8c64-871a757b9993\", \"f0a7c6ca-206e-477c-8f0a-416a355532d5\": \"f0a7c6ca-206e-477c-8f0a-416a355532d5\"}, \"doc_id_dict\": {}, \"embeddings_dict\": {}}"
-    }
-  }
-}
+{"index_store/data": {"f0aa9b38-a1b4-4cd8-8261-38f7e28b7053": {"__type__": "vector_store", "__data__": "{\"index_id\": \"f0aa9b38-a1b4-4cd8-8261-38f7e28b7053\", \"summary\": null, \"nodes_dict\": {\"432d9349-9c31-4ff4-b982-7a3d80a383a4\": \"432d9349-9c31-4ff4-b982-7a3d80a383a4\", \"1c93f814-8019-43d5-a583-12e1c1c1d702\": \"1c93f814-8019-43d5-a583-12e1c1c1d702\", \"6850a6e1-7bf9-42c8-873c-fff410bf4f61\": \"6850a6e1-7bf9-42c8-873c-fff410bf4f61\", \"99678815-82b6-4df5-88e5-8f489ef5f0dd\": \"99678815-82b6-4df5-88e5-8f489ef5f0dd\", \"f56f6064-9d14-4290-bfb5-b5992fddbfd8\": \"f56f6064-9d14-4290-bfb5-b5992fddbfd8\", \"d8cd4c80-4fcd-4961-a6f9-e315078bdd3b\": \"d8cd4c80-4fcd-4961-a6f9-e315078bdd3b\", \"3b4b0e7f-dcde-4cf2-94a5-e96fa67d52bb\": \"3b4b0e7f-dcde-4cf2-94a5-e96fa67d52bb\", \"1a275084-07c0-41cb-85d5-20b454630cde\": \"1a275084-07c0-41cb-85d5-20b454630cde\", \"0fdf0b9e-db8c-46d3-a634-67bd4b3939ca\": \"0fdf0b9e-db8c-46d3-a634-67bd4b3939ca\", \"a428ca4f-d82d-4dbd-98ef-63dce24a749d\": \"a428ca4f-d82d-4dbd-98ef-63dce24a749d\", \"1f537ab9-d15c-456a-937e-481b1b84dfd1\": \"1f537ab9-d15c-456a-937e-481b1b84dfd1\", \"598cd507-a5c5-4b6d-a555-037d1fc388d5\": \"598cd507-a5c5-4b6d-a555-037d1fc388d5\", \"65f05413-df67-4a93-80cf-83654439fdd8\": \"65f05413-df67-4a93-80cf-83654439fdd8\", \"dc8cbec7-d6e5-462c-977a-738f1eedcffc\": \"dc8cbec7-d6e5-462c-977a-738f1eedcffc\", \"a5771b3f-00f7-4669-b593-647827e570c8\": \"a5771b3f-00f7-4669-b593-647827e570c8\", \"cc945889-e48c-4262-9654-86d60bfe38b6\": \"cc945889-e48c-4262-9654-86d60bfe38b6\", \"1e6cc8f4-ec6c-474f-8db2-93de82248a28\": \"1e6cc8f4-ec6c-474f-8db2-93de82248a28\", \"d9441ba7-09ec-4726-ae88-df1863261a9f\": \"d9441ba7-09ec-4726-ae88-df1863261a9f\", \"a5aa5af4-12d1-4a01-8e0e-5318ca3b6de4\": \"a5aa5af4-12d1-4a01-8e0e-5318ca3b6de4\", \"009e44e3-4730-485f-add7-e1e517ee1536\": \"009e44e3-4730-485f-add7-e1e517ee1536\", \"3b0649c0-a990-4340-8d4d-b65db3cb432b\": \"3b0649c0-a990-4340-8d4d-b65db3cb432b\", \"bd7b998c-ead4-4d14-a01d-9c4c3d1d5d65\": \"bd7b998c-ead4-4d14-a01d-9c4c3d1d5d65\"}, \"doc_id_dict\": {}, \"embeddings_dict\": {}}"}}}

langchain_manager/manager.py

@@ -5,8 +5,6 @@ from langchain.embeddings.base import Embeddings as LCEmbeddings
 from langchain.embeddings.openai import OpenAIEmbeddings
 from langchain.llms import AzureOpenAI
 
-from core.lifecycle import Lifecycle
-
 
 class BaseLangChainManager(ABC):
     def __init__(self) -> None:
@@ -39,3 +37,29 @@ class LangChainAzureManager(BaseLangChainManager):
     # Override
     def get_llm(self) -> BaseLanguageModel:
         return self.llm
+
+
+class LangChainHuggingFaceManager(BaseLangChainManager):
+    def __init__(self) -> None:
+        super().__init__()
+        from transformers import AutoTokenizer, AutoModel
+
+        AutoTokenizer.from_pretrained("GanymedeNil/text2vec-large-chinese")
+
+        AutoModel.from_pretrained("GanymedeNil/text2vec-large-chinese")
+
+        self.embedding = OpenAIEmbeddings(client=None, chunk_size=1)
+        self.llm = AzureOpenAI(
+            deployment_name="text-davinci-003",
+            # model_name="text-davinci-003",
+            model="text-davinci-003",
+            client=None,
+        )
+
+    # Override
+    def get_embedding(self) -> LCEmbeddings:
+        return self.embedding
+
+    # Override
+    def get_llm(self) -> BaseLanguageModel:
+        return self.llm

llama/data_loader.py

@@ -51,7 +51,8 @@ class GithubLoader(WikiLoader):
             verbose=True,
             concurrent_requests=10,
         )
-
+        os.environ["http_proxy"] = "http://127.0.0.1:7890"
+        os.environ["https_proxy"] = "http://127.0.0.1:7890"
         docs = loader.load_data(branch="master")
 
         with open("docs/docs.pkl", "wb") as f:

llama/{context.py → service_context.py}

@@ -66,6 +66,65 @@ class AzureServiceContextManager(ServiceContextManager):
         )
 
 
+class HuggingFaceChineseOptServiceContextManager(ServiceContextManager):
+    lc_manager: BaseLangChainManager
+    service_context: ServiceContext
+
+    def __init__(self, lc_manager: BaseLangChainManager):
+        super().__init__()
+        self.lc_manager = lc_manager
+
+    def get_service_context(self) -> ServiceContext:
+        if self.service_context is None:
+            raise ValueError(
+                "service context is not ready, check for lifecycle statement"
+            )
+        return self.service_context
+
+    def do_init(self) -> None:
+        # define embedding
+        from langchain.embeddings import HuggingFaceEmbeddings
+
+        model_name = "GanymedeNil/text2vec-large-chinese"
+        hf_embedding = HuggingFaceEmbeddings(
+            model_name=model_name, model_kwargs={"device": "cpu"}
+        )
+
+        embedding = LangchainEmbedding(hf_embedding)
+        # define LLM
+        llm_predictor = LLMPredictor(self.lc_manager.get_llm())
+        # configure service context
+        self.service_context = ServiceContext.from_defaults(
+            llm_predictor=llm_predictor, embed_model=embedding
+        )
+
+    def do_start(self) -> None:
+        self.logger.info(
+            "[do_start][embedding] last used usage: %d",
+            self.service_context.embed_model.total_tokens_used,
+        )
+        self.logger.info(
+            "[do_start][predict] last used usage: %d",
+            self.service_context.llm_predictor.total_tokens_used,
+        )
+
+    def do_stop(self) -> None:
+        self.logger.info(
+            "[do_stop][embedding] last used usage: %d",
+            self.service_context.embed_model.total_tokens_used,
+        )
+        self.logger.info(
+            "[do_stop][predict] last used usage: %d",
+            self.service_context.llm_predictor.total_tokens_used,
+        )
+
+    def do_dispose(self) -> None:
+        self.logger.info(
+            "[do_dispose] total used token: %d",
+            self.service_context.llm_predictor.total_tokens_used,
+        )
+
+
 class StorageContextManager(Lifecycle, ABC):
     @abstractmethod
     def get_storage_context(self) -> StorageContext:

local-requirements.txt

@@ -0,0 +1 @@
+python-dotenv

requirements.txt

@@ -1,6 +1,7 @@
 llama_index>=0.6.3
 llama_hub
-
+streamlit
 ruff
 black
-mypy
+mypy
+accelerate

xpipe_wiki/manager_factory.py

@@ -4,39 +4,45 @@ import os
 from core.helper import LifecycleHelper
 from xpipe_wiki.robot_manager import XPipeWikiRobotManager, AzureXPipeWikiRobotManager
 
+from multiprocessing import Lock
 
-class XPipeRobotRevision(enum.Enum):
-    SIMPLE_OPENAI_VERSION_0 = 1
-
-
+lock = Lock()
 
 
+class XPipeRobotRevision(enum.Enum):
+    SIMPLE_OPENAI_VERSION_0 = 1
+    HUGGINGFACE_VERSION_0 = 2
 
 
 class XPipeRobotManagerFactory:
     """
-    CAPABLE: Dict[XPipeRobotRevision, XPipeWikiRobotManager] = {XPipeRobotRevision.SIMPLE_OPENAI_VERSION_0: XPipeWikiRobotManager()}
+    CAPABLE: Dict[XPipeRobotRevision, XPipeWikiRobotManager] =
+    {XPipeRobotRevision.SIMPLE_OPENAI_VERSION_0: XPipeWikiRobotManager()}
     """
-    CAPABLE = dict()
+
+    CAPABLE = dict()  # type: dict[XPipeRobotRevision, XPipeWikiRobotManager]
+
     @classmethod
     def get_or_create(cls, revision: XPipeRobotRevision) -> XPipeWikiRobotManager:
-        if cls.CAPABLE.get(revision) is not None:
-            return cls.CAPABLE[revision]
-        if revision == XPipeRobotRevision.SIMPLE_OPENAI_VERSION_0:
-            manager = cls.create_simple_openai_version_0()
-        cls.CAPABLE[revision] = manager
-        return manager
+        with lock:
+            if cls.CAPABLE.get(revision) is not None:
+                return cls.CAPABLE[revision]
+            if revision == XPipeRobotRevision.SIMPLE_OPENAI_VERSION_0:
+                manager = cls.create_simple_openai_version_0()
+            elif revision == XPipeRobotRevision.HUGGINGFACE_VERSION_0:
+                manager = cls.create_huggingface_version_0()
+            cls.CAPABLE[revision] = manager
+            return manager
 
     @classmethod
     def create_simple_openai_version_0(cls) -> AzureXPipeWikiRobotManager:
-        from llama.context import AzureServiceContextManager
+        from llama.service_context import AzureServiceContextManager
         from langchain_manager.manager import LangChainAzureManager
 
         service_context_manager = AzureServiceContextManager(
             lc_manager=LangChainAzureManager()
         )
-
-        from llama.context import LocalStorageContextManager
+        from llama.service_context import LocalStorageContextManager
 
         dataset_path = os.getenv("XPIPE_WIKI_DATASET_PATH", "./dataset")
         storage_context_manager = LocalStorageContextManager(
@@ -51,3 +57,26 @@ class XPipeRobotManagerFactory:
         LifecycleHelper.start_if_possible(robot_manager)
         return robot_manager
 
+    @classmethod
+    def create_huggingface_version_0(cls) -> AzureXPipeWikiRobotManager:
+        from llama.service_context import HuggingFaceChineseOptServiceContextManager
+        from langchain_manager.manager import LangChainAzureManager
+
+        service_context_manager = HuggingFaceChineseOptServiceContextManager(
+            lc_manager=LangChainAzureManager()
+        )
+
+        from llama.service_context import LocalStorageContextManager
+
+        dataset_path = os.getenv("XPIPE_WIKI_DATASET_PATH", "./dataset")
+        storage_context_manager = LocalStorageContextManager(
+            dataset_path=dataset_path, service_context_manager=service_context_manager
+        )
+
+        robot_manager = AzureXPipeWikiRobotManager(
+            service_context_manager=service_context_manager,
+            storage_context_manager=storage_context_manager,
+        )
+        LifecycleHelper.initialize_if_possible(robot_manager)
+        LifecycleHelper.start_if_possible(robot_manager)
+        return robot_manager

xpipe_wiki/robot_manager.py

@@ -3,11 +3,10 @@ from typing import Any
 
 from llama_index import load_index_from_storage
 from llama_index.indices.query.base import BaseQueryEngine
-from pydantic import dataclasses
 
 from core.helper import LifecycleHelper
 from core.lifecycle import Lifecycle
-from llama.context import ServiceContextManager, StorageContextManager
+from llama.service_context import ServiceContextManager, StorageContextManager
 
 
 class XPipeWikiRobot(ABC):
@@ -24,7 +23,7 @@ class AzureOpenAIXPipeWikiRobot(XPipeWikiRobot):
         self.query_engine = query_engine
 
     def ask(self, question: str) -> Any:
-        return self.query_engine.query(question)
+        return self.query_engine.query(question).response
 
 
 class XPipeWikiRobotManager(Lifecycle):
@@ -58,7 +57,8 @@ class AzureXPipeWikiRobotManager(XPipeWikiRobotManager):
         LifecycleHelper.start_if_possible(self.service_context_manager)
         LifecycleHelper.start_if_possible(self.storage_context_manager)
         index = load_index_from_storage(
-            storage_context=self.storage_context_manager.get_storage_context()
+            storage_context=self.storage_context_manager.get_storage_context(),
+            service_context=self.service_context_manager.get_service_context(),
         )
         self.query_engine = index.as_query_engine(
             service_context=self.service_context_manager.get_service_context()