Spaces:
Sleeping
Sleeping
from langchain_community.vectorstores import FAISS | |
from langchain.embeddings.huggingface import HuggingFaceEmbeddings | |
from langchain_core.output_parsers import StrOutputParser | |
from langchain_core.runnables import RunnablePassthrough | |
from langchain_core.runnables import RunnableParallel | |
import streamlit as st | |
import re | |
from langchain.llms.base import LLM | |
from langchain.llms.utils import enforce_stop_tokens | |
from typing import Dict, List, Optional, Tuple, Union | |
import requests | |
import json | |
embeddings = HuggingFaceEmbeddings(model_name='/Users/yunshi/Downloads/同步空间/LLM/2023ChatGPT/Coding/RAG/bge-large-zh/') ## 切换成BGE的embedding。 | |
# embeddings = HuggingFaceEmbeddings(model_name='/Users/yunshi/Downloads/360Data/Data Center/Working-On Task/演讲与培训/2023ChatGPT/RAG/bge-large-zh/') ## 切换成BGE的embedding。 | |
vector_store = FAISS.load_local("./faiss_index/", embeddings=embeddings, allow_dangerous_deserialization=True) ## 加载vector store到本地。 | |
# vector_store = FAISS.load_local("./faiss_index/", embeddings=embeddings) ## 加载vector store到本地。 ### original code here. | |
## 配置ChatGLM的类与后端api server对应。 | |
class ChatGLM(LLM): | |
max_token: int = 8096 ### 无法输出response的时候,可以看一下这里。 | |
temperature: float = 0.7 | |
top_p = 0.9 | |
history = [] | |
def __init__(self): | |
super().__init__() | |
def _llm_type(self) -> str: | |
return "ChatGLM" | |
def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str: | |
# headers中添加上content-type这个参数,指定为json格式 | |
headers = {'Content-Type': 'application/json'} | |
data=json.dumps({ | |
'prompt':prompt, | |
'temperature':self.temperature, | |
'history':self.history, | |
'max_length':self.max_token | |
}) | |
print("ChatGLM prompt:",prompt) | |
# 调用api | |
# response = requests.post("http://0.0.0.0:8000",headers=headers,data=data) ##working。 | |
response = requests.post("http://127.0.0.1:8000",headers=headers,data=data) ##working。 | |
print("ChatGLM resp:", response) | |
if response.status_code!=200: | |
return "查询结果错误" | |
resp = response.json() | |
if stop is not None: | |
response = enforce_stop_tokens(response, stop) | |
self.history = self.history+[[None, resp['response']]] ##original | |
return resp['response'] ##original. | |
## 在绝对路径中提取完整的文件名 | |
def extract_document_name(path): | |
# 路径分割 | |
path_segments = path.split("/") | |
# 文件名提取 | |
document_name = path_segments[-1] | |
return document_name | |
## 从一段话中提取 1 句完整的句子,且该句子的长度必须超过 5 个词,同时去除了换行符'\n\n'。 | |
import re | |
def extract_sentence(text): | |
""" | |
从一段话中提取 1 句完整的句子,且该句子的长度必须超过 5 个词。 | |
Args: | |
text: 一段话。 | |
Returns: | |
提取到的句子。 | |
""" | |
# 去除换行符。 | |
text = text.replace('\n\n', '') | |
# 使用正则表达式匹配句子。 | |
sentences = re.split(r'[。?!;]', text) | |
# 过滤掉长度小于 5 个词的句子。 | |
sentences = [sentence for sentence in sentences if len(sentence.split()) >= 5] | |
# 返回第一句句子。 | |
return sentences[0] if sentences else None | |
### 综合source的输出内容。 | |
def rag_source(docs): | |
print('starting source function!') | |
source = "" | |
for i, doc in enumerate(docs): | |
source += f"**【信息来源 {i+1}】** " + extract_document_name(doc.metadata['source']) + ',' + f"第{docs[i].metadata['page']+1}页" + ',部分内容摘录:' + extract_sentence(doc.page_content) + '\n\n' | |
print('source:', source) | |
return source | |
def rag_response(user_input, k=3): | |
# docs = vector_store.similarity_search('user_input', k=k) ## Original。 | |
docs = vector_store.similarity_search(user_input, k=k) ##TODO 'user_input' to user_input? | |
context = [doc.page_content for doc in docs] | |
# print('context: {}'.format(context)) | |
source = rag_source(docs=docs) ## 封装到一个函数中。 | |
## 用大模型来回答问题。 | |
llm = ChatGLM() ## 启动一个实例。 | |
# final_prompt = f"已知信息:\n{context}\n 根据这些已知信息来回答问题:\n{user_input}" | |
final_prompt = f"已知信息:\n{context}\n 根据这些已知信息尽可能详细且专业地来回答问题:\n{user_input}" | |
response = llm(prompt=final_prompt) ## 通过实例化之后的LLM来输出结果。 | |
# response = llm(prompt=final_prompt) ## 通过实例化之后的LLM来输出结果。 | |
# response = llm(prompt='where is shanghai') | |
# print('response now:' + response) | |
return response, source | |
# # import asyncio | |
# response, source = rag_response('我是一个企业主,我需要关注哪些存货的数据资源规则?') | |
# print(response) | |
# print(source) |