添加模型

pilot/configs/model_config.py

@@ -13,9 +13,11 @@ LOGDIR = os.path.join(root_path, "logs")
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 llm_model_config = {
     "flan-t5-base": os.path.join(model_path, "flan-t5-base"),
-    "vicuna-13b": os.path.join(model_path, "vicuna-13b")
+    "vicuna-13b": os.path.join(model_path, "vicuna-13b"),
+    "sentence-transforms": os.path.join(model_path, "all-MiniLM-L6-v2")
 }
 
+
 LLM_MODEL = "vicuna-13b"
 LIMIT_MODEL_CONCURRENCY = 5
 MAX_POSITION_EMBEDDINGS = 2048

pilot/vector_store/extract_tovec.py

@@ -5,6 +5,8 @@
 from langchain.text_splitter import CharacterTextSplitter
 from langchain.vectorstores import Chroma
 from pilot.model.vicuna_llm import VicunaEmbeddingLLM
+# from langchain.embeddings import SentenceTransformerEmbeddings
+
 
 embeddings = VicunaEmbeddingLLM()
 
@@ -20,4 +22,17 @@ def knownledge_tovec(filename):
     return docsearch
 
 
+# def knownledge_tovec_st(filename):
+#     """ Use sentence transformers to embedding the document.
+#         https://github.com/UKPLab/sentence-transformers
+#     """
+#     from pilot.configs.model_config import llm_model_config
+#     embeddings = SentenceTransformerEmbeddings(model=llm_model_config["sentence-transforms"])
 
+#     with open(filename, "r") as f:
+#         knownledge = f.read()
+    
+#     text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
+#     texts = text_splitter(knownledge)
+#     docsearch = Chroma.from_texts(texts, embeddings, metadatas=[{"source": str(i)} for i in range(len(texts))])
+#     return docsearch