File size: 1,642 Bytes
99b76c6 aed726e 99b76c6 aed726e 99b76c6 83c198a 99b76c6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
from transformers import AutoTokenizer, AutoModel
from langchain.document_loaders import TextLoader
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.llms import HuggingFaceHub
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import FAISS
import datetime
# 选择 embedding 算法
model_name = "sentence-transformers/all-mpnet-base-v2"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': False}
eb = HuggingFaceEmbeddings(
model_name=model_name,
model_kwargs=model_kwargs,
encode_kwargs=encode_kwargs
)
# 从文件读取数据
loader = TextLoader("./test.txt")
docs = loader.load()
# 分词
text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=50)
tl = []
for doc in docs:
for t in text_splitter.split_text(doc.page_content):
tl.append(t)
# 向量数据库
vector_store = FAISS.from_texts(tl, eb)
query = '为什么client不直接从data server拉取数据'
docs = vector_store.similarity_search(query, k=4)
# 构造提问
context = ''
for i in range(len(docs)):
context += f'{i+1}、{docs[i].page_content}\n'
prompt = f'已知:\n{context}\n问题:\n{query}'
print(prompt)
print(f'[{datetime.datetime.now()}]start to init model')
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b-int4", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm-6b-int4", trust_remote_code=True).float()
model = model.eval()
print(f'[{datetime.datetime.now()}]model ready, start to chat')
response, history = model.chat(tokenizer, prompt, history=[])
print(f'[{datetime.datetime.now()}]response:{response}')
|