一、和chatglm3通信，基於chromadb
import json
import requests
import os
from pdfminer.high_level import extract_pages
from pdfminer.layout import LTTextContainer
import chromadb
from chromadb.config import Settings
import numpy as np
from numpy import dot
from numpy.linalg import norm
from sentence_transformers import SentenceTransformer
###########################classes&functions########################
model = SentenceTransformer('/home/helu/milvus/m3e-base')
# 調用本地GLM接口
def get_completion_glm(prompt):
endpoint_url = "http://192.168.212.211:6006/"
chat_messages = [
{
"role": "system",
"content": "從現在開始扮演一個專業人士和我對話",
}]
chat_messages.append({"role": "user", "content": prompt})
data = {
"model": "chatglm3", # 模型名稱
"messages": chat_messages, # 會話歷史
"stream": False, # 是否流式響應
"max_tokens": 2000, # 最多生成字數
"temperature": 0.8, # 溫度
"top_p": 0.8, # 採樣概率
}
headers = {'Content-Type': 'application/json'}
#response = requests.request("POST", url, headers=headers, data=payload).json()
response = requests.post(f"{endpoint_url}v1/chat/completions", headers=headers, json=data,stream=False)
parsed_response = response.json()
text = parsed_response.get("choices", [{}])[0].get("message", "").get("content", "")
return text
def extract_text_from_pdf(filename, page_numbers=None, min_line_length=1):
'''從 PDF 文件中（按指定頁碼）提取文字'''
paragraphs = []
buffer = ''
full_text = ''
# 提取全部文本
for i, page_layout in enumerate(extract_pages(filename)):
# 如果指定了頁碼範圍，跳過範圍外的頁
if page_numbers is not None and i not in page_numbers:
continue
for element in page_layout:
if isinstance(element, LTTextContainer):
full_text += element.get_text() + '\n'
# 按空行分隔，將文本重新組織成段落
lines = full_text.split('\n')
for text in lines:
if len(text) >= min_line_length:
buffer += (' '+text) if not text.endswith('-') else text.strip('-')
elif buffer:
paragraphs.append(buffer)
buffer = ''
if buffer:
paragraphs.append(buffer)
return paragraphs
prompt_template = """
你是一個問答機器人。
你的任務是根據下述給定的已知信息回答用戶問題。
確保你的回覆完全依據下述已知信息。不要編造答案。
如果下述已知信息不足以回答用戶的問題，請直接回復"我無法回答您的問題"。
已知信息:
__INFO__
用戶問：
__QUERY__
請用中文回答用戶問題。
"""
def build_prompt(prompt_template, **kwargs):
'''將 Prompt 模板賦值'''
prompt = prompt_template
for k, v in kwargs.items():
if isinstance(v, str):
val = v
elif isinstance(v, list) and all(isinstance(elem, str) for elem in v):
val = '\n'.join(v)
else:
val = str(v)
prompt = prompt.replace(f"__{k.upper()}__", val)
return prompt
###需要換成本地接口###
def get_embeddings(texts):
#data = embedding.create(input=texts).data
embeddings = model.encode(texts)
#return [x.embedding for x in data]
return embeddings
class MyVectorDBConnector:
def __init__(self, collection_name, embedding_fn):
chroma_client = chromadb.Client(Settings(allow_reset=True))
#chroma_client.reset()
# 創建一個 collection
self.collection = chroma_client.get_or_create_collection(name=collection_name)
self.embedding_fn = embedding_fn
def add_documents(self, documents):
'''向 collection 中添加文檔與向量'''
self.collection.add(
embeddings=self.embedding_fn(documents), # 每個文檔的向量
documents=documents, # 文檔的原文
ids=[f"id{i}" for i in range(len(documents))] # 每個文檔的 id
)
def search(self, query, top_n):
'''檢索向量數據庫'''
results = self.collection.query(
query_embeddings=self.embedding_fn([query]),
n_results=top_n
)
return results
class RAG_Bot:
def __init__(self, vector_db, llm_api, n_results=2):
self.vector_db = vector_db
self.llm_api = llm_api
self.n_results = n_results
def chat(self, user_query):
# 1. 檢索
search_results = self.vector_db.search(user_query, self.n_results)
# 2. 構建 Prompt
prompt = build_prompt(
prompt_template, info=search_results['documents'][0], query=user_query)
# 3. 調用 LLM
response = self.llm_api(prompt)
return response
############################################################################
# 只取兩頁（第一章）
paragraphs = extract_text_from_pdf("llama2.pdf", page_numbers=[2, 3], min_line_length=10)
# 創建一個向量數據庫對象（注意這種寫法）
vector_db = MyVectorDBConnector("demo", get_embeddings)
# 向向量數據庫中添加文檔
vector_db.add_documents(paragraphs)
# 創建一個RAG機器人
bot = RAG_Bot(vector_db,llm_api=get_completion_glm)
user_query = "llama 2有對話版嗎？"
response = bot.chat(user_query)
print(response)
二、和chatglm3通信，基於faiss
import json
import requests
import os
from pdfminer.high_level import extract_pages
from pdfminer.layout import LTTextContainer
import faiss
import numpy as np
from numpy import dot
from numpy.linalg import norm
from sentence_transformers import SentenceTransformer
###########################classes&functions########################
model = SentenceTransformer('/home/helu/milvus/m3e-base')
# 調用對話接口
def get_completion_glm(prompt):
endpoint_url = "http://192.168.212.211:6006/" #DL4 GPU，回答10秒
#endpoint_url = "http://127.0.0.1:8000/" #筆記本虛擬機,回答2分鐘
chat_messages = [
{
"role": "system",
"content": "從現在開始扮演一個專業人士和我對話",
}]
chat_messages.append({"role": "user", "content": prompt})
data = {
"model": "chatglm3", # 模型名稱
"messages": chat_messages, # 會話歷史
"stream": False, # 是否流式響應
"max_tokens": 2000, # 最多生成字數
"temperature": 0.8, # 溫度
"top_p": 0.8, # 採樣概率
}
headers = {'Content-Type': 'application/json'}
response = requests.post(f"{endpoint_url}v1/chat/completions", headers=headers, json=data,stream=False)
parsed_response = response.json()
text = parsed_response.get("choices", [{}])[0].get("message", "").get("content", "")
return text
def extract_text_from_pdf(filename, page_numbers=None, min_line_length=1):
'''從 PDF 文件中（按指定頁碼）提取文字'''
paragraphs = []
buffer = ''
full_text = ''
# 提取全部文本
for i, page_layout in enumerate(extract_pages(filename)):
# 如果指定了頁碼範圍，跳過範圍外的頁
if page_numbers is not None and i not in page_numbers:
continue
for element in page_layout:
if isinstance(element, LTTextContainer):
full_text += element.get_text() + '\n'
# 按空行分隔，將文本重新組織成段落
lines = full_text.split('\n')
for text in lines:
if len(text) >= min_line_length:
buffer += (' '+text) if not text.endswith('-') else text.strip('-')
elif buffer:
paragraphs.append(buffer)
buffer = ''
if buffer:
paragraphs.append(buffer)
return paragraphs
prompt_template = """
你是一個問答機器人。
你的任務是根據下述給定的已知信息回答用戶問題。
確保你的回覆完全依據下述已知信息。不要編造答案。
如果下述已知信息不足以回答用戶的問題，請直接回復"我無法回答您的問題"。
已知信息:
__INFO__
用戶問：
__QUERY__
請用中文回答用戶問題。
"""
def build_prompt(prompt_template, **kwargs):
'''將 Prompt 模板賦值'''
prompt = prompt_template
for k, v in kwargs.items():
if isinstance(v, str):
val = v
elif isinstance(v, list) and all(isinstance(elem, str) for elem in v):
val = '\n'.join(v)
else:
val = str(v)
prompt = prompt.replace(f"__{k.upper()}__", val)
return prompt
def get_embeddings(texts):
embeddings = model.encode(texts)
return embeddings
def create_index(datas_embedding):
index = faiss.IndexFlatL2(datas_embedding.shape[1]) # 這裏必須傳入一個向量的維度，創建一個空的索引
index.add(datas_embedding) # 把向量數據加入索引
return index
###需要換成本地接口###
class RAG_Bot:
def __init__(self, vector_db, llm_api, paragraphs,n_results=2):
self.vector_db = vector_db
self.llm_api = llm_api
self.n_results = n_results
self.para = paragraphs
def chat(self, user_query):
# 1. 檢索
query_embedding = model.encode([user_query])
Distance, Index = self.vector_db.search(query_embedding, self.n_results)
# 2. 構建 Prompt
par = paragraphs[int(Index[0][0])]
prompt = build_prompt(
prompt_template, info=par, query=user_query)
# 3. 調用 LLM
response = self.llm_api(prompt)
return response
############################################################################
# 只取兩頁（第一章）
paragraphs = extract_text_from_pdf("llama2.pdf", page_numbers=[2, 3], min_line_length=10)
datas_embedding = get_embeddings(paragraphs)
faiss_index = create_index(datas_embedding)
# 創建一個RAG機器人
bot = RAG_Bot(faiss_index,llm_api=get_completion_glm,paragraphs = paragraphs)
user_query = "llama2相比較llama1有哪些提高？"
response = bot.chat(user_query)
print(response)
初步感覺，chromadb更接近一個數據庫，它有完整的檢索過程，可以得到原文和索引
def chat(self, user_query):
# 1. 檢索
search_results = self.vector_db.search(user_query, self.n_results)
# 2. 構建 Prompt
prompt = build_prompt(
prompt_template, info=search_results['documents'][0], query=user_query)
# 3. 調用 LLM
response = self.llm_api(prompt)
return response
而faiss更像是一個輕量的檢索工具，搜索過程中只能夠得到索引，構建prompt的過程中需要原文
def chat(self, user_query):
# 1. 檢索
query_embedding = model.encode([user_query])
Distance, Index = self.vector_db.search(query_embedding, self.n_results)
# 2. 構建 Prompt
par = paragraphs[int(Index[0][0])]
prompt = build_prompt(
prompt_template, info=par, query=user_query)
# 3. 調用 LLM
response = self.llm_api(prompt)
return response