feat(RAG): 更新RAG模块代码和文档
refactor: 简化Embeddings和LLM类实现，移除不必要依赖 docs: 更新文档内容，添加硅基流动API使用说明 chore: 更新requirements.txt依赖版本
This commit is contained in:
KMnO4-zx
@@ -146,21 +146,59 @@ class BaseEmbeddings:
|
||||
Base class for embeddings
|
||||
"""
|
||||
def __init__(self, path: str, is_api: bool) -> None:
|
||||
"""
|
||||
初始化嵌入基类
|
||||
Args:
|
||||
path (str): 模型或数据的路径
|
||||
is_api (bool): 是否使用API方式。True表示使用在线API服务，False表示使用本地模型
|
||||
"""
|
||||
self.path = path
|
||||
self.is_api = is_api
|
||||
|
||||
def get_embedding(self, text: str, model: str) -> List[float]:
|
||||
"""
|
||||
获取文本的嵌入向量表示
|
||||
Args:
|
||||
text (str): 输入文本
|
||||
model (str): 使用的模型名称
|
||||
Returns:
|
||||
List[float]: 文本的嵌入向量
|
||||
Raises:
|
||||
NotImplementedError: 该方法需要在子类中实现
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
def cosine_similarity(cls, vector1: List[float], vector2: List[float]) -> float:
|
||||
"""
|
||||
calculate cosine similarity between two vectors
|
||||
计算两个向量之间的余弦相似度
|
||||
Args:
|
||||
vector1 (List[float]): 第一个向量
|
||||
vector2 (List[float]): 第二个向量
|
||||
Returns:
|
||||
float: 两个向量的余弦相似度，范围在[-1,1]之间
|
||||
"""
|
||||
dot_product = np.dot(vector1, vector2)
|
||||
magnitude = np.linalg.norm(vector1) * np.linalg.norm(vector2)
|
||||
if not magnitude:
|
||||
return 0
|
||||
# 将输入列表转换为numpy数组，并指定数据类型为float32
|
||||
v1 = np.array(vector1, dtype=np.float32)
|
||||
v2 = np.array(vector2, dtype=np.float32)
|
||||
|
||||
# 检查向量中是否包含无穷大或NaN值
|
||||
if not np.all(np.isfinite(v1)) or not np.all(np.isfinite(v2)):
|
||||
return 0.0
|
||||
|
||||
# 计算向量的点积
|
||||
dot_product = np.dot(v1, v2)
|
||||
# 计算向量的范数（长度）
|
||||
norm_v1 = np.linalg.norm(v1)
|
||||
norm_v2 = np.linalg.norm(v2)
|
||||
|
||||
# 计算分母（两个向量范数的乘积）
|
||||
magnitude = norm_v1 * norm_v2
|
||||
# 处理分母为0的特殊情况
|
||||
if magnitude == 0:
|
||||
return 0.0
|
||||
|
||||
# 返回余弦相似度
|
||||
return dot_product / magnitude
|
||||
```
|
||||
|
||||
@@ -176,12 +214,16 @@ class OpenAIEmbedding(BaseEmbeddings):
|
||||
def __init__(self, path: str = '', is_api: bool = True) -> None:
|
||||
super().__init__(path, is_api)
|
||||
if self.is_api:
|
||||
from openai import OpenAI
|
||||
self.client = OpenAI()
|
||||
# 从环境变量中获取 硅基流动 密钥
|
||||
self.client.api_key = os.getenv("OPENAI_API_KEY")
|
||||
# 从环境变量中获取 硅基流动 的基础URL
|
||||
self.client.base_url = os.getenv("OPENAI_BASE_URL")
|
||||
|
||||
def get_embedding(self, text: str, model: str = "text-embedding-3-large") -> List[float]:
|
||||
def get_embedding(self, text: str, model: str = "BAAI/bge-m3") -> List[float]:
|
||||
"""
|
||||
此处默认使用轨迹流动的免费嵌入模型 BAAI/bge-m3
|
||||
"""
|
||||
if self.is_api:
|
||||
text = text.replace("\n", " ")
|
||||
return self.client.embeddings.create(input=[text], model=model).data[0].embedding
|
||||
@@ -189,6 +231,9 @@ class OpenAIEmbedding(BaseEmbeddings):
|
||||
raise NotImplementedError
|
||||
```
|
||||
|
||||
> 注：此处我们默认使用国内用户可访问的硅基流动大模型API服务平台。
|
||||
> 硅基流动：https://cloud.siliconflow.cn/
|
||||
|
||||
#### Step 3: 文档加载和切分
|
||||
|
||||
接下来我们来实现一个文档加载和切分的类，这个类主要用于加载文档并将其切分成文档片段。
|
||||
@@ -251,7 +296,7 @@ def get_chunk(cls, text: str, max_token_len: int = 600, cover_content: int = 150
|
||||
- `get_vector`：获取文档的向量表示。
|
||||
- `query`：根据问题检索相关文档片段。
|
||||
|
||||
完整代码可以在 ***[RAG/VectorBase.py](RAG/VectorBase.py)*** 文件中找到。
|
||||
完整代码可以在 ***[/VectorBase.py](./RAG/VectorBase.py)*** 文件中找到。
|
||||
|
||||
```python
|
||||
class VectorStore:
|
||||
@@ -302,41 +347,43 @@ class BaseModel:
|
||||
pass
|
||||
```
|
||||
|
||||
`BaseModel` 包含两个方法：`chat`和`load_model`。对于本地化运行的开源模型需要实现`load_model`，而API模型则不需要。
|
||||
|
||||
下面以 ***[InternLM2-chat-7B](https://huggingface.co/internlm/internlm2-chat-7b)*** 模型为例：
|
||||
`BaseModel` 包含两个方法：`chat`和`load_model`。对于本地化运行的开源模型需要实现`load_model`，而API模型则不需要。在此处我们还是使用国内用户可访问的硅基流动大模型API服务平台，使用API服务的好处就是用户不需要本地的计算资源，可以大大降低学习者的学习门槛。
|
||||
|
||||
```python
|
||||
class InternLMChat(BaseModel):
|
||||
def __init__(self, path: str = '') -> None:
|
||||
super().__init__(path)
|
||||
self.load_model()
|
||||
from openai import OpenAI
|
||||
|
||||
def chat(self, prompt: str, history: List = [], content: str='') -> str:
|
||||
prompt = PROMPT_TEMPLATE['InternLM_PROMPT_TEMPLATE'].format(question=prompt, context=content)
|
||||
response, history = self.model.chat(self.tokenizer, prompt, history)
|
||||
return response
|
||||
class OpenAIChat(BaseModel):
|
||||
def __init__(self, model: str = "Qwen/Qwen2.5-32B-Instruct") -> None:
|
||||
self.model = model
|
||||
|
||||
def chat(self, prompt: str, history: List[dict], content: str) -> str:
|
||||
client = OpenAI()
|
||||
client.api_key = os.getenv("OPENAI_API_KEY")
|
||||
client.base_url = os.getenv("OPENAI_BASE_URL")
|
||||
history.append({'role': 'user', 'content': RAG_PROMPT_TEMPLATE.format(question=prompt, context=content)})
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=history,
|
||||
max_tokens=2048,
|
||||
temperature=0.1
|
||||
)
|
||||
return response.choices[0].message.content
|
||||
|
||||
def load_model(self):
|
||||
import torch
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(self.path, trust_remote_code=True)
|
||||
self.model = AutoModelForCausalLM.from_pretrained(self.path, torch_dtype=torch.float16, trust_remote_code=True).cuda()
|
||||
```
|
||||
|
||||
可以用一个字典来保存所有的prompt，方便维护：
|
||||
设计一个专用于RAG的大模型提示词，如下：
|
||||
|
||||
```python
|
||||
PROMPT_TEMPLATE = dict(
|
||||
InternLM_PROMPT_TEMPLATE="""先对上下文进行内容总结,再使用上下文来回答用户的问题。如果你不知道答案，就说你不知道。总是使用中文回答。
|
||||
问题: {question}
|
||||
可参考的上下文：
|
||||
···
|
||||
{context}
|
||||
···
|
||||
如果给定的上下文无法让你做出回答，请回答数据库中没有这个内容，你不知道。
|
||||
有用的回答:"""
|
||||
)
|
||||
RAG_PROMPT_TEMPLATE="""
|
||||
使用以上下文来回答用户的问题。如果你不知道答案，就说你不知道。总是使用中文回答。
|
||||
问题: {question}
|
||||
可参考的上下文：
|
||||
···
|
||||
{context}
|
||||
···
|
||||
如果给定的上下文无法让你做出回答，请回答数据库中没有这个内容，你不知道。
|
||||
有用的回答:
|
||||
"""
|
||||
```
|
||||
|
||||
这样我们就可以利用InternLM2模型来做RAG啦！
|
||||
@@ -346,47 +393,51 @@ PROMPT_TEMPLATE = dict(
|
||||
接下来，我们来看看Tiny-RAG的Demo吧！
|
||||
|
||||
```python
|
||||
from RAG.VectorBase import VectorStore
|
||||
from RAG.utils import ReadFiles
|
||||
from RAG.LLM import OpenAIChat, InternLMChat
|
||||
from RAG.Embeddings import JinaEmbedding, ZhipuEmbedding
|
||||
from VectorBase import VectorStore
|
||||
from utils import ReadFiles
|
||||
from LLM import OpenAIChat
|
||||
from Embeddings import OpenAIEmbedding
|
||||
|
||||
# 没有保存数据库
|
||||
docs = ReadFiles('./data').get_content(max_token_len=600, cover_content=150) # 获取data目录下的所有文件内容并分割
|
||||
docs = ReadFiles('./data').get_content(max_token_len=600, cover_content=150) # 获得data目录下的所有文件内容并分割
|
||||
vector = VectorStore(docs)
|
||||
embedding = ZhipuEmbedding() # 创建EmbeddingModel
|
||||
embedding = OpenAIEmbedding() # 创建EmbeddingModel
|
||||
vector.get_vector(EmbeddingModel=embedding)
|
||||
vector.persist(path='storage') # 将向量和文档内容保存到storage目录，下次再用可以直接加载本地数据库
|
||||
vector.persist(path='storage') # 将向量和文档内容保存到storage目录下，下次再用就可以直接加载本地的数据库
|
||||
|
||||
question = 'git的原理是什么？'
|
||||
# vector.load_vector('./storage') # 加载本地的数据库
|
||||
|
||||
content = vector.query(question, model='zhipu', k=1)[0]
|
||||
chat = InternLMChat(path='model_path')
|
||||
question = 'RAG的原理是什么？'
|
||||
|
||||
content = vector.query(question, EmbeddingModel=embedding, k=1)[0]
|
||||
chat = OpenAIChat(model='Qwen/Qwen2.5-32B-Instruct')
|
||||
print(chat.chat(question, [], content))
|
||||
```
|
||||
|
||||
也可以从本地加载已处理好的数据库：
|
||||
|
||||
```python
|
||||
from RAG.VectorBase import VectorStore
|
||||
from RAG.utils import ReadFiles
|
||||
from RAG.LLM import OpenAIChat, InternLMChat
|
||||
from RAG.Embeddings import JinaEmbedding, ZhipuEmbedding
|
||||
from VectorBase import VectorStore
|
||||
from utils import ReadFiles
|
||||
from LLM import OpenAIChat
|
||||
from Embeddings import OpenAIEmbedding
|
||||
|
||||
# 保存数据库之后
|
||||
vector = VectorStore()
|
||||
|
||||
vector.load_vector('./storage') # 加载本地数据库
|
||||
vector.load_vector('./storage') # 加载本地的数据库
|
||||
|
||||
question = 'git的原理是什么？'
|
||||
question = 'RAG的原理是什么？'
|
||||
|
||||
embedding = ZhipuEmbedding() # 创建EmbeddingModel
|
||||
|
||||
content = vector.query(question, EmbeddingModel=embedding, k=1)[0]
|
||||
chat = InternLMChat(path='model_path')
|
||||
chat = OpenAIChat(model='Qwen/Qwen2.5-32B-Instruct')
|
||||
print(chat.chat(question, [], content))
|
||||
```
|
||||
|
||||
> 注：7.2 章节的所有代码均可在 [Happy-LLM Chapter7 RAG](https://github.com/datawhalechina/happy-llm/tree/main/docs/chapter7/RAG) 中找到。
|
||||
|
||||
## 7.3 Agent
|
||||
|
||||
### 7.3.1 什么是 LLM Agent？
|
||||
|
||||
Reference in New Issue
Block a user