github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-07-23 12:21:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

添加模型

Browse Source
This commit is contained in:
csunny 2023-05-04 15:00:21 +08:00
parent 3d9bdfa216
commit ffeb5040cc
2 changed files with 18 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
pilot/configs/model_config.py
View File

@ -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

15
pilot/vector_store/extract_tovec.py
View File

@ -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