Initial commit
This commit is contained in:
@@ -0,0 +1,75 @@
|
||||
import os
|
||||
# hugging face镜像设置，如果国内环境无法使用启用该设置
|
||||
# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
|
||||
from dotenv import load_dotenv
|
||||
from langchain_community.document_loaders import UnstructuredMarkdownLoader
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
from langchain_huggingface import HuggingFaceEmbeddings
|
||||
from langchain_core.vectorstores import InMemoryVectorStore
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_openai import ChatOpenAI
|
||||
|
||||
load_dotenv()
|
||||
|
||||
markdown_path = "../../data/C1/markdown/easy-rl-chapter1.md"
|
||||
|
||||
# 加载本地markdown文件
|
||||
loader = UnstructuredMarkdownLoader(markdown_path)
|
||||
docs = loader.load()
|
||||
|
||||
# 文本分块
|
||||
text_splitter = RecursiveCharacterTextSplitter()
|
||||
chunks = text_splitter.split_documents(docs)
|
||||
|
||||
# 中文嵌入模型
|
||||
embeddings = HuggingFaceEmbeddings(
|
||||
model_name="BAAI/bge-small-zh-v1.5",
|
||||
model_kwargs={'device': 'cpu'},
|
||||
encode_kwargs={'normalize_embeddings': True}
|
||||
)
|
||||
|
||||
# 构建向量存储
|
||||
vectorstore = InMemoryVectorStore(embeddings)
|
||||
vectorstore.add_documents(chunks)
|
||||
|
||||
# 提示词模板
|
||||
prompt = ChatPromptTemplate.from_template("""请根据下面提供的上下文信息来回答问题。
|
||||
请确保你的回答完全基于这些上下文。
|
||||
如果上下文中没有足够的信息来回答问题，请直接告知：“抱歉，我无法根据提供的上下文找到相关信息来回答此问题。”
|
||||
|
||||
上下文:
|
||||
{context}
|
||||
|
||||
问题: {question}
|
||||
|
||||
回答:"""
|
||||
)
|
||||
|
||||
# 配置大语言模型
|
||||
|
||||
# 使用 AIHubmix
|
||||
llm = ChatOpenAI(
|
||||
model="glm-4.7-flash-free",
|
||||
temperature=0.7,
|
||||
max_tokens=4096,
|
||||
api_key=os.getenv("DEEPSEEK_API_KEY"),
|
||||
base_url="https://aihubmix.com/v1"
|
||||
)
|
||||
|
||||
# llm = ChatOpenAI(
|
||||
# model="deepseek-chat",
|
||||
# temperature=0.7,
|
||||
# max_tokens=4096,
|
||||
# api_key=os.getenv("DEEPSEEK_API_KEY"),
|
||||
# base_url="https://api.deepseek.com"
|
||||
# )
|
||||
|
||||
# 用户查询
|
||||
question = "文中举了哪些例子？"
|
||||
|
||||
# 在向量存储中查询相关文档
|
||||
retrieved_docs = vectorstore.similarity_search(question, k=3)
|
||||
docs_content = "\n\n".join(doc.page_content for doc in retrieved_docs)
|
||||
|
||||
answer = llm.invoke(prompt.format(question=question, context=docs_content))
|
||||
print(answer)
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
# os.environ['HF_ENDPOINT']='https://hf-mirror.com'
|
||||
from dotenv import load_dotenv
|
||||
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
|
||||
from llama_index.llms.openai_like import OpenAILike
|
||||
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# 使用 AIHubmix
|
||||
Settings.llm = OpenAILike(
|
||||
model="glm-4.7-flash-free",
|
||||
api_key=os.getenv("DEEPSEEK_API_KEY"),
|
||||
api_base="https://aihubmix.com/v1",
|
||||
is_chat_model=True
|
||||
)
|
||||
|
||||
# Settings.llm = OpenAI(
|
||||
# model="deepseek-chat",
|
||||
# api_key=os.getenv("DEEPSEEK_API_KEY"),
|
||||
# api_base="https://api.deepseek.com"
|
||||
# )
|
||||
Settings.embed_model = HuggingFaceEmbedding("BAAI/bge-small-zh-v1.5")
|
||||
|
||||
docs = SimpleDirectoryReader(input_files=["../../data/C1/markdown/easy-rl-chapter1.md"]).load_data()
|
||||
|
||||
index = VectorStoreIndex.from_documents(docs)
|
||||
|
||||
query_engine = index.as_query_engine()
|
||||
|
||||
print(query_engine.get_prompts())
|
||||
|
||||
print(query_engine.query("文中举了哪些例子?"))
|
||||
@@ -0,0 +1,4 @@
|
||||
import nltk
|
||||
|
||||
nltk.download('punkt', force=True)
|
||||
nltk.download('averaged_perceptron_tagger', force=True)
|
||||
@@ -0,0 +1,25 @@
|
||||
from unstructured.partition.auto import partition
|
||||
|
||||
# PDF文件路径
|
||||
pdf_path = "../../data/C2/pdf/rag.pdf"
|
||||
|
||||
# 使用Unstructured加载并解析PDF文档
|
||||
elements = partition(
|
||||
filename=pdf_path,
|
||||
content_type="application/pdf"
|
||||
)
|
||||
|
||||
# 打印解析结果
|
||||
print(f"解析完成: {len(elements)} 个元素, {sum(len(str(e)) for e in elements)} 字符")
|
||||
|
||||
# 统计元素类型
|
||||
from collections import Counter
|
||||
types = Counter(e.category for e in elements)
|
||||
print(f"元素类型: {dict(types)}")
|
||||
|
||||
# 显示所有元素
|
||||
print("\n所有元素:")
|
||||
for i, element in enumerate(elements, 1):
|
||||
print(f"Element {i} ({element.category}):")
|
||||
print(element)
|
||||
print("=" * 60)
|
||||
@@ -0,0 +1,23 @@
|
||||
from langchain.text_splitter import CharacterTextSplitter
|
||||
from langchain_community.document_loaders import TextLoader
|
||||
|
||||
# 1. 文档加载
|
||||
loader = TextLoader("../../data/C2/txt/蜂医.txt", encoding="utf-8")
|
||||
docs = loader.load()
|
||||
|
||||
# 2. 初始化固定大小分块器
|
||||
text_splitter = CharacterTextSplitter(
|
||||
chunk_size=200, # 每个块的大小
|
||||
chunk_overlap=10 # 块之间的重叠大小
|
||||
)
|
||||
|
||||
# 3. 执行分块
|
||||
chunks = text_splitter.split_documents(docs)
|
||||
|
||||
# 4. 打印结果
|
||||
print(f"文本被切分为 {len(chunks)} 个块。\n")
|
||||
print("--- 前5个块内容示例 ---")
|
||||
for i, chunk in enumerate(chunks[:5]):
|
||||
print("=" * 60)
|
||||
# chunk 是一个 Document 对象，需要访问它的 .page_content 属性来获取文本
|
||||
print(f'块 {i+1} (长度: {len(chunk.page_content)}): "{chunk.page_content}"')
|
||||
@@ -0,0 +1,20 @@
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
from langchain_community.document_loaders import TextLoader
|
||||
|
||||
loader = TextLoader("../../data/C2/txt/蜂医.txt", encoding="utf-8")
|
||||
docs = loader.load()
|
||||
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
# 针对中英文混合文本，定义一个更全面的分隔符列表
|
||||
separators=["\n\n", "\n", "。", "，", " ", ""], # 按顺序尝试分割
|
||||
chunk_size=200,
|
||||
chunk_overlap=10
|
||||
)
|
||||
|
||||
chunks = text_splitter.split_documents(docs)
|
||||
|
||||
print(f"文本被切分为 {len(chunks)} 个块。\n")
|
||||
print("--- 前5个块内容示例 ---")
|
||||
for i, chunk in enumerate(chunks[:5]):
|
||||
print("=" * 60)
|
||||
print(f'块 {i+1} (长度: {len(chunk.page_content)}): "{chunk.page_content}"')
|
||||
@@ -0,0 +1,26 @@
|
||||
from langchain_experimental.text_splitter import SemanticChunker
|
||||
from langchain_community.embeddings import HuggingFaceEmbeddings
|
||||
from langchain_community.document_loaders import TextLoader
|
||||
|
||||
embeddings = HuggingFaceEmbeddings(
|
||||
model_name="BAAI/bge-small-zh-v1.5",
|
||||
model_kwargs={'device': 'cpu'},
|
||||
encode_kwargs={'normalize_embeddings': True}
|
||||
)
|
||||
|
||||
# 初始化 SemanticChunker
|
||||
text_splitter = SemanticChunker(
|
||||
embeddings,
|
||||
breakpoint_threshold_type="percentile" # 也可以是 "standard_deviation", "interquartile", "gradient"
|
||||
)
|
||||
|
||||
loader = TextLoader("../../data/C2/txt/蜂医.txt", encoding="utf-8")
|
||||
documents = loader.load()
|
||||
|
||||
docs = text_splitter.split_documents(documents)
|
||||
|
||||
print(f"文本被切分为 {len(docs)} 个块。\n")
|
||||
print("--- 前2个块内容示例 ---")
|
||||
for i, chunk in enumerate(docs[:2]):
|
||||
print("=" * 60)
|
||||
print(f'块 {i+1} (长度: {len(chunk.page_content)}):\n"{chunk.page_content}"')
|
||||
@@ -0,0 +1,32 @@
|
||||
import torch
|
||||
from visual_bge.visual_bge.modeling import Visualized_BGE
|
||||
|
||||
model = Visualized_BGE(model_name_bge="BAAI/bge-base-en-v1.5",
|
||||
model_weight="../../models/bge/Visualized_base_en_v1.5.pth")
|
||||
model.eval()
|
||||
|
||||
with torch.no_grad():
|
||||
text_emb = model.encode(text="datawhale开源组织的logo")
|
||||
img_emb_1 = model.encode(image="../../data/C3/imgs/datawhale01.png")
|
||||
multi_emb_1 = model.encode(image="../../data/C3/imgs/datawhale01.png", text="datawhale开源组织的logo")
|
||||
img_emb_2 = model.encode(image="../../data/C3/imgs/datawhale02.png")
|
||||
multi_emb_2 = model.encode(image="../../data/C3/imgs/datawhale02.png", text="datawhale开源组织的logo")
|
||||
|
||||
# 计算相似度
|
||||
sim_1 = img_emb_1 @ img_emb_2.T
|
||||
sim_2 = img_emb_1 @ multi_emb_1.T
|
||||
sim_3 = text_emb @ multi_emb_1.T
|
||||
sim_4 = multi_emb_1 @ multi_emb_2.T
|
||||
|
||||
print("=== 相似度计算结果 ===")
|
||||
print(f"纯图像 vs 纯图像: {sim_1}")
|
||||
print(f"图文结合1 vs 纯图像: {sim_2}")
|
||||
print(f"图文结合1 vs 纯文本: {sim_3}")
|
||||
print(f"图文结合1 vs 图文结合2: {sim_4}")
|
||||
|
||||
# 向量信息分析
|
||||
print("\n=== 嵌入向量信息 ===")
|
||||
print(f"多模态向量维度: {multi_emb_1.shape}")
|
||||
print(f"图像向量维度: {img_emb_1.shape}")
|
||||
print(f"多模态向量示例 (前10个元素): {multi_emb_1[0][:10]}")
|
||||
print(f"图像向量示例 (前10个元素): {img_emb_1[0][:10]}")
|
||||
@@ -0,0 +1,37 @@
|
||||
from langchain_community.vectorstores import FAISS
|
||||
from langchain_community.embeddings import HuggingFaceEmbeddings
|
||||
from langchain_core.documents import Document
|
||||
|
||||
# 1. 示例文本和嵌入模型
|
||||
texts = [
|
||||
"张三是法外狂徒",
|
||||
"FAISS是一个用于高效相似性搜索和密集向量聚类的库。",
|
||||
"LangChain是一个用于开发由语言模型驱动的应用程序的框架。"
|
||||
]
|
||||
docs = [Document(page_content=t) for t in texts]
|
||||
embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-zh-v1.5")
|
||||
|
||||
# 2. 创建向量存储并保存到本地
|
||||
vectorstore = FAISS.from_documents(docs, embeddings)
|
||||
|
||||
local_faiss_path = "./faiss_index_store"
|
||||
vectorstore.save_local(local_faiss_path)
|
||||
|
||||
print(f"FAISS index has been saved to {local_faiss_path}")
|
||||
|
||||
# 3. 加载索引并执行查询
|
||||
# 加载时需指定相同的嵌入模型，并允许反序列化
|
||||
loaded_vectorstore = FAISS.load_local(
|
||||
local_faiss_path,
|
||||
embeddings,
|
||||
allow_dangerous_deserialization=True
|
||||
)
|
||||
|
||||
# 执行相似性搜索
|
||||
query = "FAISS是做什么的？"
|
||||
results = loaded_vectorstore.similarity_search(query, k=1)
|
||||
|
||||
print(f"\n查询: '{query}'")
|
||||
print("相似度最高的文档:")
|
||||
for doc in results:
|
||||
print(f"- {doc.page_content}")
|
||||
@@ -0,0 +1,19 @@
|
||||
from llama_index.core import VectorStoreIndex, Document, Settings
|
||||
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
|
||||
|
||||
# 1. 配置全局嵌入模型
|
||||
Settings.embed_model = HuggingFaceEmbedding("BAAI/bge-small-zh-v1.5")
|
||||
|
||||
# 2. 创建示例文档
|
||||
texts = [
|
||||
"张三是法外狂徒",
|
||||
"LlamaIndex是一个用于构建和查询私有或领域特定数据的框架。",
|
||||
"它提供了数据连接、索引和查询接口等工具。"
|
||||
]
|
||||
docs = [Document(text=t) for t in texts]
|
||||
|
||||
# 3. 创建索引并持久化到本地
|
||||
index = VectorStoreIndex.from_documents(docs)
|
||||
persist_path = "./llamaindex_index_store"
|
||||
index.storage_context.persist(persist_dir=persist_path)
|
||||
print(f"LlamaIndex 索引已保存至: {persist_path}")
|
||||
@@ -0,0 +1,160 @@
|
||||
import os
|
||||
from tqdm import tqdm
|
||||
from glob import glob
|
||||
import torch
|
||||
from visual_bge.visual_bge.modeling import Visualized_BGE
|
||||
from pymilvus import MilvusClient, FieldSchema, CollectionSchema, DataType
|
||||
import numpy as np
|
||||
import cv2
|
||||
from PIL import Image
|
||||
|
||||
# 1. 初始化设置
|
||||
MODEL_NAME = "BAAI/bge-base-en-v1.5"
|
||||
MODEL_PATH = "../../models/bge/Visualized_base_en_v1.5.pth"
|
||||
DATA_DIR = "../../data/C3"
|
||||
COLLECTION_NAME = "multimodal_demo"
|
||||
MILVUS_URI = "http://localhost:19530"
|
||||
|
||||
# 2. 定义工具 (编码器和可视化函数)
|
||||
class Encoder:
|
||||
"""编码器类，用于将图像和文本编码为向量。"""
|
||||
def __init__(self, model_name: str, model_path: str):
|
||||
self.model = Visualized_BGE(model_name_bge=model_name, model_weight=model_path)
|
||||
self.model.eval()
|
||||
|
||||
def encode_query(self, image_path: str, text: str) -> list[float]:
|
||||
with torch.no_grad():
|
||||
query_emb = self.model.encode(image=image_path, text=text)
|
||||
return query_emb.tolist()[0]
|
||||
|
||||
def encode_image(self, image_path: str) -> list[float]:
|
||||
with torch.no_grad():
|
||||
query_emb = self.model.encode(image=image_path)
|
||||
return query_emb.tolist()[0]
|
||||
|
||||
def visualize_results(query_image_path: str, retrieved_images: list, img_height: int = 300, img_width: int = 300, row_count: int = 3) -> np.ndarray:
|
||||
"""从检索到的图像列表创建一个全景图用于可视化。"""
|
||||
panoramic_width = img_width * row_count
|
||||
panoramic_height = img_height * row_count
|
||||
panoramic_image = np.full((panoramic_height, panoramic_width, 3), 255, dtype=np.uint8)
|
||||
query_display_area = np.full((panoramic_height, img_width, 3), 255, dtype=np.uint8)
|
||||
|
||||
# 处理查询图像
|
||||
query_pil = Image.open(query_image_path).convert("RGB")
|
||||
query_cv = np.array(query_pil)[:, :, ::-1]
|
||||
resized_query = cv2.resize(query_cv, (img_width, img_height))
|
||||
bordered_query = cv2.copyMakeBorder(resized_query, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=(255, 0, 0))
|
||||
query_display_area[img_height * (row_count - 1):, :] = cv2.resize(bordered_query, (img_width, img_height))
|
||||
cv2.putText(query_display_area, "Query", (10, panoramic_height - 20), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
|
||||
|
||||
# 处理检索到的图像
|
||||
for i, img_path in enumerate(retrieved_images):
|
||||
row, col = i // row_count, i % row_count
|
||||
start_row, start_col = row * img_height, col * img_width
|
||||
|
||||
retrieved_pil = Image.open(img_path).convert("RGB")
|
||||
retrieved_cv = np.array(retrieved_pil)[:, :, ::-1]
|
||||
resized_retrieved = cv2.resize(retrieved_cv, (img_width - 4, img_height - 4))
|
||||
bordered_retrieved = cv2.copyMakeBorder(resized_retrieved, 2, 2, 2, 2, cv2.BORDER_CONSTANT, value=(0, 0, 0))
|
||||
panoramic_image[start_row:start_row + img_height, start_col:start_col + img_width] = bordered_retrieved
|
||||
|
||||
# 添加索引号
|
||||
cv2.putText(panoramic_image, str(i), (start_col + 10, start_row + 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
|
||||
|
||||
return np.hstack([query_display_area, panoramic_image])
|
||||
|
||||
# 3. 初始化客户端
|
||||
print("--> 正在初始化编码器和Milvus客户端...")
|
||||
encoder = Encoder(MODEL_NAME, MODEL_PATH)
|
||||
milvus_client = MilvusClient(uri=MILVUS_URI)
|
||||
|
||||
# 4. 创建 Milvus Collection
|
||||
print(f"\n--> 正在创建 Collection '{COLLECTION_NAME}'")
|
||||
if milvus_client.has_collection(COLLECTION_NAME):
|
||||
milvus_client.drop_collection(COLLECTION_NAME)
|
||||
print(f"已删除已存在的 Collection: '{COLLECTION_NAME}'")
|
||||
|
||||
image_list = glob(os.path.join(DATA_DIR, "dragon", "*.png"))
|
||||
if not image_list:
|
||||
raise FileNotFoundError(f"在 {DATA_DIR}/dragon/ 中未找到任何 .png 图像。")
|
||||
dim = len(encoder.encode_image(image_list[0]))
|
||||
|
||||
fields = [
|
||||
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
|
||||
FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
||||
FieldSchema(name="image_path", dtype=DataType.VARCHAR, max_length=512),
|
||||
]
|
||||
|
||||
# 创建集合 Schema
|
||||
schema = CollectionSchema(fields, description="多模态图文检索")
|
||||
print("Schema 结构:")
|
||||
print(schema)
|
||||
|
||||
# 创建集合
|
||||
milvus_client.create_collection(collection_name=COLLECTION_NAME, schema=schema)
|
||||
print(f"成功创建 Collection: '{COLLECTION_NAME}'")
|
||||
print("Collection 结构:")
|
||||
print(milvus_client.describe_collection(collection_name=COLLECTION_NAME))
|
||||
|
||||
# 5. 准备并插入数据
|
||||
print(f"\n--> 正在向 '{COLLECTION_NAME}' 插入数据")
|
||||
data_to_insert = []
|
||||
for image_path in tqdm(image_list, desc="生成图像嵌入"):
|
||||
vector = encoder.encode_image(image_path)
|
||||
data_to_insert.append({"vector": vector, "image_path": image_path})
|
||||
|
||||
if data_to_insert:
|
||||
result = milvus_client.insert(collection_name=COLLECTION_NAME, data=data_to_insert)
|
||||
print(f"成功插入 {result['insert_count']} 条数据。")
|
||||
|
||||
# 6. 创建索引
|
||||
print(f"\n--> 正在为 '{COLLECTION_NAME}' 创建索引")
|
||||
index_params = milvus_client.prepare_index_params()
|
||||
index_params.add_index(
|
||||
field_name="vector",
|
||||
index_type="HNSW",
|
||||
metric_type="COSINE",
|
||||
params={"M": 16, "efConstruction": 256}
|
||||
)
|
||||
milvus_client.create_index(collection_name=COLLECTION_NAME, index_params=index_params)
|
||||
print("成功为向量字段创建 HNSW 索引。")
|
||||
print("索引详情:")
|
||||
print(milvus_client.describe_index(collection_name=COLLECTION_NAME, index_name="vector"))
|
||||
milvus_client.load_collection(collection_name=COLLECTION_NAME)
|
||||
print("已加载 Collection 到内存中。")
|
||||
|
||||
# 7. 执行多模态检索
|
||||
print(f"\n--> 正在 '{COLLECTION_NAME}' 中执行检索")
|
||||
query_image_path = os.path.join(DATA_DIR, "dragon", "query.png")
|
||||
query_text = "一条龙"
|
||||
query_vector = encoder.encode_query(image_path=query_image_path, text=query_text)
|
||||
|
||||
search_results = milvus_client.search(
|
||||
collection_name=COLLECTION_NAME,
|
||||
data=[query_vector],
|
||||
output_fields=["image_path"],
|
||||
limit=5,
|
||||
search_params={"metric_type": "COSINE", "params": {"ef": 128}}
|
||||
)[0]
|
||||
|
||||
retrieved_images = []
|
||||
print("检索结果:")
|
||||
for i, hit in enumerate(search_results):
|
||||
print(f" Top {i+1}: ID={hit['id']}, 距离={hit['distance']:.4f}, 路径='{hit['entity']['image_path']}'")
|
||||
retrieved_images.append(hit['entity']['image_path'])
|
||||
|
||||
# 8. 可视化与清理
|
||||
print(f"\n--> 正在可视化结果并清理资源")
|
||||
if not retrieved_images:
|
||||
print("没有检索到任何图像。")
|
||||
else:
|
||||
panoramic_image = visualize_results(query_image_path, retrieved_images)
|
||||
combined_image_path = os.path.join(DATA_DIR, "search_result.png")
|
||||
cv2.imwrite(combined_image_path, panoramic_image)
|
||||
print(f"结果图像已保存到: {combined_image_path}")
|
||||
Image.open(combined_image_path).show()
|
||||
|
||||
milvus_client.release_collection(collection_name=COLLECTION_NAME)
|
||||
print(f"已从内存中释放 Collection: '{COLLECTION_NAME}'")
|
||||
milvus_client.drop_collection(COLLECTION_NAME)
|
||||
print(f"已删除 Collection: '{COLLECTION_NAME}'")
|
||||
@@ -0,0 +1,51 @@
|
||||
import os
|
||||
from llama_index.core.node_parser import SentenceWindowNodeParser, SentenceSplitter
|
||||
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
|
||||
from llama_index.llms.deepseek import DeepSeek
|
||||
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
|
||||
from llama_index.core.postprocessor import MetadataReplacementPostProcessor
|
||||
|
||||
# 1. 配置模型
|
||||
Settings.llm = DeepSeek(model="deepseek-chat", temperature=0.1, api_key=os.getenv("DEEPSEEK_API_KEY"))
|
||||
Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en")
|
||||
|
||||
# 2. 加载文档
|
||||
documents = SimpleDirectoryReader(
|
||||
input_files=["../../data/C3/pdf/IPCC_AR6_WGII_Chapter03.pdf"]
|
||||
).load_data()
|
||||
|
||||
# 3. 创建节点与构建索引
|
||||
# 3.1 句子窗口索引
|
||||
node_parser = SentenceWindowNodeParser.from_defaults(
|
||||
window_size=3,
|
||||
window_metadata_key="window",
|
||||
original_text_metadata_key="original_text",
|
||||
)
|
||||
sentence_nodes = node_parser.get_nodes_from_documents(documents)
|
||||
sentence_index = VectorStoreIndex(sentence_nodes)
|
||||
|
||||
# 3.2 常规分块索引 (基准)
|
||||
base_parser = SentenceSplitter(chunk_size=512)
|
||||
base_nodes = base_parser.get_nodes_from_documents(documents)
|
||||
base_index = VectorStoreIndex(base_nodes)
|
||||
|
||||
# 4. 构建查询引擎
|
||||
sentence_query_engine = sentence_index.as_query_engine(
|
||||
similarity_top_k=2,
|
||||
node_postprocessors=[
|
||||
MetadataReplacementPostProcessor(target_metadata_key="window")
|
||||
],
|
||||
)
|
||||
base_query_engine = base_index.as_query_engine(similarity_top_k=2)
|
||||
|
||||
# 5. 执行查询并对比结果
|
||||
query = "What are the concerns surrounding the AMOC?"
|
||||
print(f"查询: {query}\n")
|
||||
|
||||
print("--- 句子窗口检索结果 ---")
|
||||
window_response = sentence_query_engine.query(query)
|
||||
print(f"回答: {window_response}\n")
|
||||
|
||||
print("--- 常规检索结果 ---")
|
||||
base_response = base_query_engine.query(query)
|
||||
print(f"回答: {base_response}\n")
|
||||
@@ -0,0 +1,63 @@
|
||||
import os
|
||||
import pandas as pd
|
||||
from dotenv import load_dotenv
|
||||
from llama_index.core import VectorStoreIndex
|
||||
from llama_index.core.schema import IndexNode
|
||||
from llama_index.experimental.query_engine import PandasQueryEngine
|
||||
from llama_index.core.retrievers import RecursiveRetriever
|
||||
from llama_index.core.query_engine import RetrieverQueryEngine
|
||||
from llama_index.llms.deepseek import DeepSeek
|
||||
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
|
||||
from llama_index.core import Settings
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# 配置模型
|
||||
Settings.llm = DeepSeek(model="deepseek-chat", api_key=os.getenv("DEEPSEEK_API_KEY"))
|
||||
Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-zh-v1.5")
|
||||
|
||||
# 1.加载数据并为每个工作表创建查询引擎和摘要节点
|
||||
excel_file = '../../data/C3/excel/movie.xlsx'
|
||||
xls = pd.ExcelFile(excel_file)
|
||||
|
||||
df_query_engines = {}
|
||||
all_nodes = []
|
||||
|
||||
for sheet_name in xls.sheet_names:
|
||||
df = pd.read_excel(xls, sheet_name=sheet_name)
|
||||
|
||||
# 为当前工作表（DataFrame）创建一个 PandasQueryEngine
|
||||
query_engine = PandasQueryEngine(df=df, llm=Settings.llm, verbose=True)
|
||||
|
||||
# 为当前工作表创建一个摘要节点（IndexNode）
|
||||
year = sheet_name.replace('年份_', '')
|
||||
summary = f"这个表格包含了年份为 {year} 的电影信息，可以用来回答关于这一年电影的具体问题。"
|
||||
node = IndexNode(text=summary, index_id=sheet_name)
|
||||
all_nodes.append(node)
|
||||
|
||||
# 存储工作表名称到其查询引擎的映射
|
||||
df_query_engines[sheet_name] = query_engine
|
||||
|
||||
# 2. 创建顶层索引（只包含摘要节点）
|
||||
vector_index = VectorStoreIndex(all_nodes)
|
||||
|
||||
# 3. 创建递归检索器
|
||||
# 3.1 创建顶层检索器，用于在摘要节点中检索
|
||||
vector_retriever = vector_index.as_retriever(similarity_top_k=1)
|
||||
|
||||
# 3.2 创建递归检索器
|
||||
recursive_retriever = RecursiveRetriever(
|
||||
"vector",
|
||||
retriever_dict={"vector": vector_retriever},
|
||||
query_engine_dict=df_query_engines,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
# 4. 创建查询引擎
|
||||
query_engine = RetrieverQueryEngine.from_args(recursive_retriever)
|
||||
|
||||
# 5. 执行查询
|
||||
query = "1994年评分人数最少的电影是哪一部？"
|
||||
print(f"查询: {query}")
|
||||
response = query_engine.query(query)
|
||||
print(f"回答: {response}")
|
||||
@@ -0,0 +1,101 @@
|
||||
import os
|
||||
import pandas as pd
|
||||
from dotenv import load_dotenv
|
||||
from llama_index.core import VectorStoreIndex, Document, Settings
|
||||
from llama_index.core.retrievers import VectorIndexRetriever
|
||||
from llama_index.core.query_engine import RetrieverQueryEngine
|
||||
from llama_index.core.vector_stores import MetadataFilters, ExactMatchFilter
|
||||
from llama_index.llms.deepseek import DeepSeek
|
||||
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# 配置模型
|
||||
Settings.llm = DeepSeek(model="deepseek-chat", api_key=os.getenv("DEEPSEEK_API_KEY"))
|
||||
Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-zh-v1.5")
|
||||
|
||||
# 1. 加载和预处理数据
|
||||
excel_file = '../../data/C3/excel/movie.xlsx'
|
||||
xls = pd.ExcelFile(excel_file)
|
||||
|
||||
summary_docs = []
|
||||
content_docs = []
|
||||
|
||||
print("开始加载和处理Excel文件...")
|
||||
for sheet_name in xls.sheet_names:
|
||||
df = pd.read_excel(xls, sheet_name=sheet_name)
|
||||
|
||||
# 数据清洗
|
||||
if '评分人数' in df.columns:
|
||||
df['评分人数'] = df['评分人数'].astype(str).str.replace('人评价', '').str.strip()
|
||||
df['评分人数'] = pd.to_numeric(df['评分人数'], errors='coerce').fillna(0).astype(int)
|
||||
|
||||
# 创建摘要文档 (用于路由)
|
||||
year = sheet_name.replace('年份_', '')
|
||||
summary_text = f"这个表格包含了年份为 {year} 的电影信息，包括电影名称、导演、评分、评分人数等。"
|
||||
summary_doc = Document(
|
||||
text=summary_text,
|
||||
metadata={"sheet_name": sheet_name}
|
||||
)
|
||||
summary_docs.append(summary_doc)
|
||||
|
||||
# 创建内容文档 (用于最终问答)
|
||||
content_text = df.to_string(index=False)
|
||||
content_doc = Document(
|
||||
text=content_text,
|
||||
metadata={"sheet_name": sheet_name}
|
||||
)
|
||||
content_docs.append(content_doc)
|
||||
|
||||
print("数据加载和处理完成。\n")
|
||||
|
||||
# 2. 构建向量索引
|
||||
# 使用默认的内存SimpleVectorStore，它支持元数据过滤
|
||||
|
||||
# 2.1 为摘要创建索引
|
||||
summary_index = VectorStoreIndex(summary_docs)
|
||||
|
||||
# 2.2 为内容创建索引
|
||||
content_index = VectorStoreIndex(content_docs)
|
||||
|
||||
print("摘要索引和内容索引构建完成。\n")
|
||||
|
||||
# 3. 定义两步式查询逻辑
|
||||
def query_safe_recursive(query_str):
|
||||
print(f"--- 开始执行查询 ---")
|
||||
print(f"查询: {query_str}")
|
||||
|
||||
# 第一步：路由 - 在摘要索引中找到最相关的表格
|
||||
print("\n第一步：在摘要索引中进行路由...")
|
||||
summary_retriever = VectorIndexRetriever(index=summary_index, similarity_top_k=1)
|
||||
retrieved_nodes = summary_retriever.retrieve(query_str)
|
||||
|
||||
if not retrieved_nodes:
|
||||
return "抱歉，未能找到相关的电影年份信息。"
|
||||
|
||||
# 获取匹配到的工作表名称
|
||||
matched_sheet_name = retrieved_nodes[0].node.metadata['sheet_name']
|
||||
print(f"路由结果：匹配到工作表 -> {matched_sheet_name}")
|
||||
|
||||
# 第二步：检索 - 在内容索引中根据工作表名称过滤并检索具体内容
|
||||
print("\n第二步：在内容索引中检索具体信息...")
|
||||
content_retriever = VectorIndexRetriever(
|
||||
index=content_index,
|
||||
similarity_top_k=1, # 通常只返回最匹配的整个表格即可
|
||||
filters=MetadataFilters(
|
||||
filters=[ExactMatchFilter(key="sheet_name", value=matched_sheet_name)]
|
||||
)
|
||||
)
|
||||
|
||||
# 创建查询引擎并执行查询
|
||||
query_engine = RetrieverQueryEngine.from_args(content_retriever)
|
||||
response = query_engine.query(query_str)
|
||||
|
||||
print("--- 查询执行结束 ---\n")
|
||||
return response
|
||||
|
||||
# 4. 执行查询
|
||||
query = "1994年评分人数最少的电影是哪一部？"
|
||||
response = query_safe_recursive(query)
|
||||
|
||||
print(f"最终回答: {response}")
|
||||
@@ -0,0 +1,68 @@
|
||||
import requests
|
||||
from pathlib import Path
|
||||
|
||||
def download_visualized_bge_model():
|
||||
"""
|
||||
下载 Visual BGE 模型权重文件
|
||||
如果模型文件不存在，则从 Hugging Face 下载
|
||||
"""
|
||||
# 定义模型路径和下载URL
|
||||
model_dir = Path("../../models/bge")
|
||||
model_file = model_dir / "Visualized_base_en_v1.5.pth"
|
||||
download_url = "https://huggingface.co/BAAI/bge-visualized/resolve/main/Visualized_base_en_v1.5.pth?download=true"
|
||||
|
||||
# 检查模型文件是否已存在
|
||||
if model_file.exists():
|
||||
print(f"模型文件已存在: {model_file}")
|
||||
print(f"文件大小: {model_file.stat().st_size / (1024*1024):.1f} MB")
|
||||
return str(model_file)
|
||||
|
||||
# 创建目录
|
||||
model_dir.mkdir(parents=True, exist_ok=True)
|
||||
print(f"创建模型目录: {model_dir}")
|
||||
|
||||
# 下载模型
|
||||
print(f"开始下载模型...")
|
||||
print(f"下载地址: {download_url}")
|
||||
|
||||
try:
|
||||
response = requests.get(download_url, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# 获取文件大小
|
||||
total_size = int(response.headers.get('content-length', 0))
|
||||
downloaded_size = 0
|
||||
|
||||
with open(model_file, 'wb') as f:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
if chunk:
|
||||
f.write(chunk)
|
||||
downloaded_size += len(chunk)
|
||||
|
||||
# 显示下载进度
|
||||
if total_size > 0:
|
||||
progress = (downloaded_size / total_size) * 100
|
||||
print(f"\r下载进度: {progress:.1f}% ({downloaded_size/(1024*1024):.1f}/{total_size/(1024*1024):.1f} MB)", end='')
|
||||
|
||||
print(f"\n模型下载完成: {model_file}")
|
||||
print(f"文件大小: {model_file.stat().st_size / (1024*1024):.1f} MB")
|
||||
return str(model_file)
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
print(f"下载失败: {e}")
|
||||
# 如果下载失败，删除不完整的文件
|
||||
if model_file.exists():
|
||||
model_file.unlink()
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"发生错误: {e}")
|
||||
if model_file.exists():
|
||||
model_file.unlink()
|
||||
return None
|
||||
|
||||
if __name__ == "__main__":
|
||||
model_path = download_visualized_bge_model()
|
||||
if model_path:
|
||||
print(f"✅ 模型准备就绪: {model_path}")
|
||||
else:
|
||||
print("❌ 模型下载失败")
|
||||
@@ -0,0 +1 @@
|
||||
from .visual_bge.modeling import Visualized_BGE
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 150 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 101 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 120 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 50 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 880 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 149 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 102 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 176 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 103 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 123 KiB |
@@ -0,0 +1,18 @@
|
||||
from setuptools import setup, find_packages
|
||||
|
||||
setup(
|
||||
name="visual_bge",
|
||||
version="0.1.0",
|
||||
description='visual_bge',
|
||||
long_description="./README.md",
|
||||
long_description_content_type="text/markdown",
|
||||
url='https://github.com/FlagOpen/FlagEmbedding/tree/master/research/visual_bge',
|
||||
packages=find_packages(),
|
||||
install_requires=[
|
||||
'torchvision',
|
||||
'timm',
|
||||
'einops',
|
||||
'ftfy'
|
||||
],
|
||||
python_requires='>=3.6',
|
||||
)
|
||||
@@ -0,0 +1,11 @@
|
||||
from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
|
||||
from .factory import create_model, create_model_and_transforms, create_model_from_pretrained, get_tokenizer, create_eva_vision_and_transforms
|
||||
from .factory import list_models, add_model_config, get_model_config, load_checkpoint
|
||||
from .loss import ClipLoss
|
||||
from .model import CLIP, CustomCLIP, CLIPTextCfg, CLIPVisionCfg,\
|
||||
convert_weights_to_lp, convert_weights_to_fp16, trace_model, get_cast_dtype
|
||||
from .openai import load_openai_model, list_openai_models
|
||||
from .pretrained import list_pretrained, list_pretrained_models_by_tag, list_pretrained_tags_by_model,\
|
||||
get_pretrained_url, download_pretrained_from_url, is_pretrained_cfg, get_pretrained_cfg, download_pretrained
|
||||
from .tokenizer import SimpleTokenizer, tokenize
|
||||
from .transform import image_transform
|
||||
Binary file not shown.
@@ -0,0 +1,2 @@
|
||||
OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
|
||||
OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)
|
||||
@@ -0,0 +1,532 @@
|
||||
# --------------------------------------------------------
|
||||
# Adapted from https://github.com/microsoft/unilm/tree/master/beit
|
||||
# --------------------------------------------------------
|
||||
import math
|
||||
import os
|
||||
from functools import partial
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
try:
|
||||
from timm.models.layers import drop_path, to_2tuple, trunc_normal_
|
||||
except:
|
||||
from timm.layers import drop_path, to_2tuple, trunc_normal_
|
||||
|
||||
from .transformer import PatchDropout
|
||||
from .rope import VisionRotaryEmbedding, VisionRotaryEmbeddingFast
|
||||
|
||||
if os.getenv('ENV_TYPE') == 'deepspeed':
|
||||
try:
|
||||
from deepspeed.runtime.activation_checkpointing.checkpointing import checkpoint
|
||||
except:
|
||||
from torch.utils.checkpoint import checkpoint
|
||||
else:
|
||||
from torch.utils.checkpoint import checkpoint
|
||||
|
||||
try:
|
||||
import xformers.ops as xops
|
||||
except ImportError:
|
||||
xops = None
|
||||
# print("Please 'pip install xformers'")
|
||||
|
||||
|
||||
class DropPath(nn.Module):
|
||||
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
|
||||
"""
|
||||
def __init__(self, drop_prob=None):
|
||||
super(DropPath, self).__init__()
|
||||
self.drop_prob = drop_prob
|
||||
|
||||
def forward(self, x):
|
||||
return drop_path(x, self.drop_prob, self.training)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
return 'p={}'.format(self.drop_prob)
|
||||
|
||||
|
||||
class Mlp(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_features,
|
||||
hidden_features=None,
|
||||
out_features=None,
|
||||
act_layer=nn.GELU,
|
||||
norm_layer=nn.LayerNorm,
|
||||
drop=0.,
|
||||
subln=False,
|
||||
|
||||
):
|
||||
super().__init__()
|
||||
out_features = out_features or in_features
|
||||
hidden_features = hidden_features or in_features
|
||||
self.fc1 = nn.Linear(in_features, hidden_features)
|
||||
self.act = act_layer()
|
||||
|
||||
self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
|
||||
|
||||
self.fc2 = nn.Linear(hidden_features, out_features)
|
||||
self.drop = nn.Dropout(drop)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.fc1(x)
|
||||
x = self.act(x)
|
||||
# x = self.drop(x)
|
||||
# commit this for the orignal BERT implement
|
||||
x = self.ffn_ln(x)
|
||||
|
||||
x = self.fc2(x)
|
||||
x = self.drop(x)
|
||||
return x
|
||||
|
||||
class SwiGLU(nn.Module):
|
||||
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0.,
|
||||
norm_layer=nn.LayerNorm, subln=False):
|
||||
super().__init__()
|
||||
out_features = out_features or in_features
|
||||
hidden_features = hidden_features or in_features
|
||||
|
||||
self.w1 = nn.Linear(in_features, hidden_features)
|
||||
self.w2 = nn.Linear(in_features, hidden_features)
|
||||
|
||||
self.act = act_layer()
|
||||
self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
|
||||
self.w3 = nn.Linear(hidden_features, out_features)
|
||||
|
||||
self.drop = nn.Dropout(drop)
|
||||
|
||||
def forward(self, x):
|
||||
x1 = self.w1(x)
|
||||
x2 = self.w2(x)
|
||||
hidden = self.act(x1) * x2
|
||||
x = self.ffn_ln(hidden)
|
||||
x = self.w3(x)
|
||||
x = self.drop(x)
|
||||
return x
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(
|
||||
self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,
|
||||
proj_drop=0., window_size=None, attn_head_dim=None, xattn=False, rope=None, subln=False, norm_layer=nn.LayerNorm):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
head_dim = dim // num_heads
|
||||
if attn_head_dim is not None:
|
||||
head_dim = attn_head_dim
|
||||
all_head_dim = head_dim * self.num_heads
|
||||
self.scale = qk_scale or head_dim ** -0.5
|
||||
|
||||
self.subln = subln
|
||||
if self.subln:
|
||||
self.q_proj = nn.Linear(dim, all_head_dim, bias=False)
|
||||
self.k_proj = nn.Linear(dim, all_head_dim, bias=False)
|
||||
self.v_proj = nn.Linear(dim, all_head_dim, bias=False)
|
||||
else:
|
||||
self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
|
||||
|
||||
if qkv_bias:
|
||||
self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
|
||||
self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
|
||||
else:
|
||||
self.q_bias = None
|
||||
self.v_bias = None
|
||||
|
||||
if window_size:
|
||||
self.window_size = window_size
|
||||
self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
|
||||
self.relative_position_bias_table = nn.Parameter(
|
||||
torch.zeros(self.num_relative_distance, num_heads)) # 2*Wh-1 * 2*Ww-1, nH
|
||||
# cls to token & token 2 cls & cls to cls
|
||||
|
||||
# get pair-wise relative position index for each token inside the window
|
||||
coords_h = torch.arange(window_size[0])
|
||||
coords_w = torch.arange(window_size[1])
|
||||
coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww
|
||||
coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
|
||||
relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
|
||||
relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
|
||||
relative_coords[:, :, 0] += window_size[0] - 1 # shift to start from 0
|
||||
relative_coords[:, :, 1] += window_size[1] - 1
|
||||
relative_coords[:, :, 0] *= 2 * window_size[1] - 1
|
||||
relative_position_index = \
|
||||
torch.zeros(size=(window_size[0] * window_size[1] + 1, ) * 2, dtype=relative_coords.dtype)
|
||||
relative_position_index[1:, 1:] = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
|
||||
relative_position_index[0, 0:] = self.num_relative_distance - 3
|
||||
relative_position_index[0:, 0] = self.num_relative_distance - 2
|
||||
relative_position_index[0, 0] = self.num_relative_distance - 1
|
||||
|
||||
self.register_buffer("relative_position_index", relative_position_index)
|
||||
else:
|
||||
self.window_size = None
|
||||
self.relative_position_bias_table = None
|
||||
self.relative_position_index = None
|
||||
|
||||
self.attn_drop = nn.Dropout(attn_drop)
|
||||
self.inner_attn_ln = norm_layer(all_head_dim) if subln else nn.Identity()
|
||||
# self.proj = nn.Linear(all_head_dim, all_head_dim)
|
||||
self.proj = nn.Linear(all_head_dim, dim)
|
||||
self.proj_drop = nn.Dropout(proj_drop)
|
||||
self.xattn = xattn
|
||||
self.xattn_drop = attn_drop
|
||||
|
||||
self.rope = rope
|
||||
|
||||
def forward(self, x, rel_pos_bias=None, attn_mask=None):
|
||||
B, N, C = x.shape
|
||||
if self.subln:
|
||||
q = F.linear(input=x, weight=self.q_proj.weight, bias=self.q_bias)
|
||||
k = F.linear(input=x, weight=self.k_proj.weight, bias=None)
|
||||
v = F.linear(input=x, weight=self.v_proj.weight, bias=self.v_bias)
|
||||
|
||||
q = q.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3) # B, num_heads, N, C
|
||||
k = k.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
|
||||
v = v.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
|
||||
else:
|
||||
|
||||
qkv_bias = None
|
||||
if self.q_bias is not None:
|
||||
qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
|
||||
|
||||
qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
|
||||
qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) # 3, B, num_heads, N, C
|
||||
q, k, v = qkv[0], qkv[1], qkv[2]
|
||||
|
||||
if self.rope:
|
||||
# slightly fast impl
|
||||
q_t = q[:, :, 1:, :]
|
||||
ro_q_t = self.rope(q_t)
|
||||
q = torch.cat((q[:, :, :1, :], ro_q_t), -2).type_as(v)
|
||||
|
||||
k_t = k[:, :, 1:, :]
|
||||
ro_k_t = self.rope(k_t)
|
||||
k = torch.cat((k[:, :, :1, :], ro_k_t), -2).type_as(v)
|
||||
|
||||
if xops is not None:
|
||||
q = q.permute(0, 2, 1, 3) # B, num_heads, N, C -> B, N, num_heads, C
|
||||
k = k.permute(0, 2, 1, 3)
|
||||
v = v.permute(0, 2, 1, 3)
|
||||
|
||||
x = xops.memory_efficient_attention(
|
||||
q, k, v,
|
||||
p=self.xattn_drop,
|
||||
scale=self.scale,
|
||||
)
|
||||
x = x.reshape(B, N, -1)
|
||||
x = self.inner_attn_ln(x)
|
||||
x = self.proj(x)
|
||||
x = self.proj_drop(x)
|
||||
else:
|
||||
q = q * self.scale
|
||||
attn = (q @ k.transpose(-2, -1))
|
||||
|
||||
if self.relative_position_bias_table is not None:
|
||||
relative_position_bias = \
|
||||
self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
|
||||
self.window_size[0] * self.window_size[1] + 1,
|
||||
self.window_size[0] * self.window_size[1] + 1, -1) # Wh*Ww,Wh*Ww,nH
|
||||
relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
|
||||
attn = attn + relative_position_bias.unsqueeze(0).type_as(attn)
|
||||
|
||||
if rel_pos_bias is not None:
|
||||
attn = attn + rel_pos_bias.type_as(attn)
|
||||
|
||||
if attn_mask is not None:
|
||||
attn_mask = attn_mask.bool()
|
||||
attn = attn.masked_fill(~attn_mask[:, None, None, :], float("-inf"))
|
||||
|
||||
attn = attn.softmax(dim=-1)
|
||||
attn = self.attn_drop(attn)
|
||||
|
||||
x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
|
||||
x = self.inner_attn_ln(x)
|
||||
x = self.proj(x)
|
||||
x = self.proj_drop(x)
|
||||
return x
|
||||
|
||||
|
||||
class Block(nn.Module):
|
||||
|
||||
def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
|
||||
drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
|
||||
window_size=None, attn_head_dim=None, xattn=False, rope=None, postnorm=False,
|
||||
subln=False, naiveswiglu=False):
|
||||
super().__init__()
|
||||
self.norm1 = norm_layer(dim)
|
||||
self.attn = Attention(
|
||||
dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
|
||||
attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim,
|
||||
xattn=xattn, rope=rope, subln=subln, norm_layer=norm_layer)
|
||||
# NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
|
||||
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
|
||||
self.norm2 = norm_layer(dim)
|
||||
mlp_hidden_dim = int(dim * mlp_ratio)
|
||||
|
||||
if naiveswiglu:
|
||||
self.mlp = SwiGLU(
|
||||
in_features=dim,
|
||||
hidden_features=mlp_hidden_dim,
|
||||
subln=subln,
|
||||
norm_layer=norm_layer,
|
||||
)
|
||||
else:
|
||||
self.mlp = Mlp(
|
||||
in_features=dim,
|
||||
hidden_features=mlp_hidden_dim,
|
||||
act_layer=act_layer,
|
||||
subln=subln,
|
||||
drop=drop
|
||||
)
|
||||
|
||||
if init_values is not None and init_values > 0:
|
||||
self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
|
||||
self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
|
||||
else:
|
||||
self.gamma_1, self.gamma_2 = None, None
|
||||
|
||||
self.postnorm = postnorm
|
||||
|
||||
def forward(self, x, rel_pos_bias=None, attn_mask=None):
|
||||
if self.gamma_1 is None:
|
||||
if self.postnorm:
|
||||
x = x + self.drop_path(self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
|
||||
x = x + self.drop_path(self.norm2(self.mlp(x)))
|
||||
else:
|
||||
x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
|
||||
x = x + self.drop_path(self.mlp(self.norm2(x)))
|
||||
else:
|
||||
if self.postnorm:
|
||||
x = x + self.drop_path(self.gamma_1 * self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
|
||||
x = x + self.drop_path(self.gamma_2 * self.norm2(self.mlp(x)))
|
||||
else:
|
||||
x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
|
||||
x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
|
||||
return x
|
||||
|
||||
|
||||
class PatchEmbed(nn.Module):
|
||||
""" Image to Patch Embedding
|
||||
"""
|
||||
def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
|
||||
super().__init__()
|
||||
img_size = to_2tuple(img_size)
|
||||
patch_size = to_2tuple(patch_size)
|
||||
num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
|
||||
self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
|
||||
self.img_size = img_size
|
||||
self.patch_size = patch_size
|
||||
self.num_patches = num_patches
|
||||
|
||||
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
|
||||
|
||||
def forward(self, x, **kwargs):
|
||||
B, C, H, W = x.shape
|
||||
# FIXME look at relaxing size constraints
|
||||
assert H == self.img_size[0] and W == self.img_size[1], \
|
||||
f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
|
||||
x = self.proj(x).flatten(2).transpose(1, 2) # [10, 3, 224, 224] -> [10, 196, 768]
|
||||
return x
|
||||
|
||||
|
||||
class RelativePositionBias(nn.Module):
|
||||
|
||||
def __init__(self, window_size, num_heads):
|
||||
super().__init__()
|
||||
self.window_size = window_size
|
||||
self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
|
||||
self.relative_position_bias_table = nn.Parameter(
|
||||
torch.zeros(self.num_relative_distance, num_heads)) # 2*Wh-1 * 2*Ww-1, nH
|
||||
# cls to token & token 2 cls & cls to cls
|
||||
|
||||
# get pair-wise relative position index for each token inside the window
|
||||
coords_h = torch.arange(window_size[0])
|
||||
coords_w = torch.arange(window_size[1])
|
||||
coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww
|
||||
coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
|
||||
relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
|
||||
relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
|
||||
relative_coords[:, :, 0] += window_size[0] - 1 # shift to start from 0
|
||||
relative_coords[:, :, 1] += window_size[1] - 1
|
||||
relative_coords[:, :, 0] *= 2 * window_size[1] - 1
|
||||
relative_position_index = \
|
||||
torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype)
|
||||
relative_position_index[1:, 1:] = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
|
||||
relative_position_index[0, 0:] = self.num_relative_distance - 3
|
||||
relative_position_index[0:, 0] = self.num_relative_distance - 2
|
||||
relative_position_index[0, 0] = self.num_relative_distance - 1
|
||||
|
||||
self.register_buffer("relative_position_index", relative_position_index)
|
||||
|
||||
def forward(self):
|
||||
relative_position_bias = \
|
||||
self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
|
||||
self.window_size[0] * self.window_size[1] + 1,
|
||||
self.window_size[0] * self.window_size[1] + 1, -1) # Wh*Ww,Wh*Ww,nH
|
||||
return relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
|
||||
|
||||
|
||||
class EVAVisionTransformer(nn.Module):
|
||||
""" Vision Transformer with support for patch or hybrid CNN input stage
|
||||
"""
|
||||
def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
|
||||
num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
|
||||
drop_path_rate=0., norm_layer=nn.LayerNorm, init_values=None, patch_dropout=0.,
|
||||
use_abs_pos_emb=True, use_rel_pos_bias=False, use_shared_rel_pos_bias=False, rope=False,
|
||||
use_mean_pooling=True, init_scale=0.001, grad_checkpointing=False, xattn=False, postnorm=False,
|
||||
pt_hw_seq_len=16, intp_freq=False, naiveswiglu=False, subln=False):
|
||||
super().__init__()
|
||||
self.image_size = img_size
|
||||
self.num_classes = num_classes
|
||||
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
|
||||
|
||||
self.patch_embed = PatchEmbed(
|
||||
img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
|
||||
num_patches = self.patch_embed.num_patches
|
||||
|
||||
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
|
||||
# self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
|
||||
if use_abs_pos_emb:
|
||||
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
|
||||
else:
|
||||
self.pos_embed = None
|
||||
self.pos_drop = nn.Dropout(p=drop_rate)
|
||||
|
||||
if use_shared_rel_pos_bias:
|
||||
self.rel_pos_bias = RelativePositionBias(window_size=self.patch_embed.patch_shape, num_heads=num_heads)
|
||||
else:
|
||||
self.rel_pos_bias = None
|
||||
|
||||
if rope:
|
||||
half_head_dim = embed_dim // num_heads // 2
|
||||
hw_seq_len = img_size // patch_size
|
||||
self.rope = VisionRotaryEmbeddingFast(
|
||||
dim=half_head_dim,
|
||||
pt_seq_len=pt_hw_seq_len,
|
||||
ft_seq_len=hw_seq_len if intp_freq else None,
|
||||
# patch_dropout=patch_dropout
|
||||
)
|
||||
else:
|
||||
self.rope = None
|
||||
|
||||
self.naiveswiglu = naiveswiglu
|
||||
|
||||
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
|
||||
self.use_rel_pos_bias = use_rel_pos_bias
|
||||
self.blocks = nn.ModuleList([
|
||||
Block(
|
||||
dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
|
||||
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
|
||||
init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None,
|
||||
xattn=xattn, rope=self.rope, postnorm=postnorm, subln=subln, naiveswiglu=naiveswiglu)
|
||||
for i in range(depth)])
|
||||
self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
|
||||
self.fc_norm = norm_layer(embed_dim) if use_mean_pooling else None
|
||||
self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
|
||||
|
||||
if self.pos_embed is not None:
|
||||
trunc_normal_(self.pos_embed, std=.02)
|
||||
|
||||
trunc_normal_(self.cls_token, std=.02)
|
||||
# trunc_normal_(self.mask_token, std=.02)
|
||||
|
||||
self.apply(self._init_weights)
|
||||
self.fix_init_weight()
|
||||
|
||||
if isinstance(self.head, nn.Linear):
|
||||
trunc_normal_(self.head.weight, std=.02)
|
||||
self.head.weight.data.mul_(init_scale)
|
||||
self.head.bias.data.mul_(init_scale)
|
||||
|
||||
# setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
|
||||
self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()
|
||||
|
||||
self.grad_checkpointing = grad_checkpointing
|
||||
|
||||
def fix_init_weight(self):
|
||||
def rescale(param, layer_id):
|
||||
param.div_(math.sqrt(2.0 * layer_id))
|
||||
|
||||
for layer_id, layer in enumerate(self.blocks):
|
||||
rescale(layer.attn.proj.weight.data, layer_id + 1)
|
||||
if self.naiveswiglu:
|
||||
rescale(layer.mlp.w3.weight.data, layer_id + 1)
|
||||
else:
|
||||
rescale(layer.mlp.fc2.weight.data, layer_id + 1)
|
||||
|
||||
def get_cast_dtype(self) -> torch.dtype:
|
||||
return self.blocks[0].mlp.fc2.weight.dtype
|
||||
|
||||
def _init_weights(self, m):
|
||||
if isinstance(m, nn.Linear):
|
||||
trunc_normal_(m.weight, std=.02)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.LayerNorm):
|
||||
nn.init.constant_(m.bias, 0)
|
||||
nn.init.constant_(m.weight, 1.0)
|
||||
|
||||
def get_num_layers(self):
|
||||
return len(self.blocks)
|
||||
|
||||
def lock(self, unlocked_groups=0, freeze_bn_stats=False):
|
||||
assert unlocked_groups == 0, 'partial locking not currently supported for this model'
|
||||
for param in self.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
@torch.jit.ignore
|
||||
def set_grad_checkpointing(self, enable=True):
|
||||
self.grad_checkpointing = enable
|
||||
|
||||
@torch.jit.ignore
|
||||
def no_weight_decay(self):
|
||||
return {'pos_embed', 'cls_token'}
|
||||
|
||||
def get_classifier(self):
|
||||
return self.head
|
||||
|
||||
def reset_classifier(self, num_classes, global_pool=''):
|
||||
self.num_classes = num_classes
|
||||
self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
|
||||
|
||||
def forward_features(self, x, return_all_features=False):
|
||||
|
||||
x = self.patch_embed(x)
|
||||
batch_size, seq_len, _ = x.size()
|
||||
|
||||
cls_tokens = self.cls_token.expand(batch_size, -1, -1) # stole cls_tokens impl from Phil Wang, thanks
|
||||
x = torch.cat((cls_tokens, x), dim=1)
|
||||
if self.pos_embed is not None:
|
||||
x = x + self.pos_embed
|
||||
x = self.pos_drop(x)
|
||||
|
||||
# a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
|
||||
if os.getenv('RoPE') == '1':
|
||||
if self.training and not isinstance(self.patch_dropout, nn.Identity):
|
||||
x, patch_indices_keep = self.patch_dropout(x)
|
||||
self.rope.forward = partial(self.rope.forward, patch_indices_keep=patch_indices_keep)
|
||||
else:
|
||||
self.rope.forward = partial(self.rope.forward, patch_indices_keep=None)
|
||||
x = self.patch_dropout(x)
|
||||
else:
|
||||
x = self.patch_dropout(x)
|
||||
|
||||
rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None
|
||||
for blk in self.blocks:
|
||||
if self.grad_checkpointing:
|
||||
# x = checkpoint(blk, x, (rel_pos_bias,))
|
||||
x = checkpoint(blk, x, rel_pos_bias)
|
||||
else:
|
||||
x = blk(x, rel_pos_bias=rel_pos_bias)
|
||||
|
||||
if not return_all_features:
|
||||
x = self.norm(x)
|
||||
if self.fc_norm is not None:
|
||||
return self.fc_norm(x.mean(1))
|
||||
else:
|
||||
return x[:, 0]
|
||||
return x
|
||||
|
||||
def forward(self, x, return_all_features=True):
|
||||
if return_all_features:
|
||||
return self.forward_features(x, return_all_features)
|
||||
x = self.forward_features(x)
|
||||
x = self.head(x)
|
||||
return x
|
||||
@@ -0,0 +1,519 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import pathlib
|
||||
import re
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
from typing import Optional, Tuple, Union, Dict, Any
|
||||
import torch
|
||||
|
||||
from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
|
||||
from .model import CLIP, CustomCLIP, convert_weights_to_lp, convert_to_custom_text_state_dict,\
|
||||
get_cast_dtype
|
||||
from .openai import load_openai_model
|
||||
from .pretrained import is_pretrained_cfg, get_pretrained_cfg, download_pretrained, list_pretrained_tags_by_model
|
||||
from .transform import image_transform
|
||||
from .tokenizer import HFTokenizer, tokenize
|
||||
from .utils import resize_clip_pos_embed, resize_evaclip_pos_embed, resize_visual_pos_embed, resize_eva_pos_embed
|
||||
|
||||
|
||||
_MODEL_CONFIG_PATHS = [Path(__file__).parent / f"model_configs/"]
|
||||
_MODEL_CONFIGS = {} # directory (model_name: config) of model architecture configs
|
||||
|
||||
|
||||
def _natural_key(string_):
|
||||
return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())]
|
||||
|
||||
|
||||
def _rescan_model_configs():
|
||||
global _MODEL_CONFIGS
|
||||
|
||||
config_ext = ('.json',)
|
||||
config_files = []
|
||||
for config_path in _MODEL_CONFIG_PATHS:
|
||||
if config_path.is_file() and config_path.suffix in config_ext:
|
||||
config_files.append(config_path)
|
||||
elif config_path.is_dir():
|
||||
for ext in config_ext:
|
||||
config_files.extend(config_path.glob(f'*{ext}'))
|
||||
|
||||
for cf in config_files:
|
||||
with open(cf, "r", encoding="utf8") as f:
|
||||
model_cfg = json.load(f)
|
||||
if all(a in model_cfg for a in ('embed_dim', 'vision_cfg', 'text_cfg')):
|
||||
_MODEL_CONFIGS[cf.stem] = model_cfg
|
||||
|
||||
_MODEL_CONFIGS = dict(sorted(_MODEL_CONFIGS.items(), key=lambda x: _natural_key(x[0])))
|
||||
|
||||
|
||||
_rescan_model_configs() # initial populate of model config registry
|
||||
|
||||
|
||||
def list_models():
|
||||
""" enumerate available model architectures based on config files """
|
||||
return list(_MODEL_CONFIGS.keys())
|
||||
|
||||
|
||||
def add_model_config(path):
|
||||
""" add model config path or file and update registry """
|
||||
if not isinstance(path, Path):
|
||||
path = Path(path)
|
||||
_MODEL_CONFIG_PATHS.append(path)
|
||||
_rescan_model_configs()
|
||||
|
||||
|
||||
def get_model_config(model_name):
|
||||
if model_name in _MODEL_CONFIGS:
|
||||
return deepcopy(_MODEL_CONFIGS[model_name])
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def get_tokenizer(model_name):
|
||||
config = get_model_config(model_name)
|
||||
tokenizer = HFTokenizer(config['text_cfg']['hf_tokenizer_name']) if 'hf_tokenizer_name' in config['text_cfg'] else tokenize
|
||||
return tokenizer
|
||||
|
||||
|
||||
# loading openai CLIP weights when is_openai=True for training
|
||||
def load_state_dict(checkpoint_path: str, map_location: str='cpu', model_key: str='model|module|state_dict', is_openai: bool=False, skip_list: list=[]):
|
||||
if is_openai:
|
||||
model = torch.jit.load(checkpoint_path, map_location="cpu").eval()
|
||||
state_dict = model.state_dict()
|
||||
for key in ["input_resolution", "context_length", "vocab_size"]:
|
||||
state_dict.pop(key, None)
|
||||
else:
|
||||
checkpoint = torch.load(checkpoint_path, map_location=map_location)
|
||||
for mk in model_key.split('|'):
|
||||
if isinstance(checkpoint, dict) and mk in checkpoint:
|
||||
state_dict = checkpoint[mk]
|
||||
break
|
||||
else:
|
||||
state_dict = checkpoint
|
||||
if next(iter(state_dict.items()))[0].startswith('module'):
|
||||
state_dict = {k[7:]: v for k, v in state_dict.items()}
|
||||
|
||||
for k in skip_list:
|
||||
if k in list(state_dict.keys()):
|
||||
logging.info(f"Removing key {k} from pretrained checkpoint")
|
||||
del state_dict[k]
|
||||
|
||||
if os.getenv('RoPE') == '1':
|
||||
for k in list(state_dict.keys()):
|
||||
if 'freqs_cos' in k or 'freqs_sin' in k:
|
||||
del state_dict[k]
|
||||
return state_dict
|
||||
|
||||
|
||||
|
||||
def load_checkpoint(model, checkpoint_path, model_key="model|module|state_dict", strict=True):
|
||||
state_dict = load_state_dict(checkpoint_path, model_key=model_key, is_openai=False)
|
||||
# detect old format and make compatible with new format
|
||||
if 'positional_embedding' in state_dict and not hasattr(model, 'positional_embedding'):
|
||||
state_dict = convert_to_custom_text_state_dict(state_dict)
|
||||
if 'text.logit_scale' in state_dict and hasattr(model, 'logit_scale'):
|
||||
state_dict['logit_scale'] = state_dict['text.logit_scale']
|
||||
del state_dict['text.logit_scale']
|
||||
|
||||
# resize_clip_pos_embed for CLIP and open CLIP
|
||||
if 'visual.positional_embedding' in state_dict:
|
||||
resize_clip_pos_embed(state_dict, model)
|
||||
# specified to eva_vit_model
|
||||
elif 'visual.pos_embed' in state_dict:
|
||||
resize_evaclip_pos_embed(state_dict, model)
|
||||
|
||||
# resize_clip_pos_embed(state_dict, model)
|
||||
incompatible_keys = model.load_state_dict(state_dict, strict=strict)
|
||||
logging.info(f"incompatible_keys.missing_keys: {incompatible_keys.missing_keys}")
|
||||
return incompatible_keys
|
||||
|
||||
def load_clip_visual_state_dict(checkpoint_path: str, map_location: str='cpu', is_openai: bool=False, skip_list:list=[]):
|
||||
state_dict = load_state_dict(checkpoint_path, map_location=map_location, is_openai=is_openai, skip_list=skip_list)
|
||||
|
||||
for k in list(state_dict.keys()):
|
||||
if not k.startswith('visual.'):
|
||||
del state_dict[k]
|
||||
for k in list(state_dict.keys()):
|
||||
if k.startswith('visual.'):
|
||||
new_k = k[7:]
|
||||
state_dict[new_k] = state_dict[k]
|
||||
del state_dict[k]
|
||||
return state_dict
|
||||
|
||||
def load_clip_text_state_dict(checkpoint_path: str, map_location: str='cpu', is_openai: bool=False, skip_list:list=[]):
|
||||
state_dict = load_state_dict(checkpoint_path, map_location=map_location, is_openai=is_openai, skip_list=skip_list)
|
||||
|
||||
for k in list(state_dict.keys()):
|
||||
if k.startswith('visual.'):
|
||||
del state_dict[k]
|
||||
return state_dict
|
||||
|
||||
def get_pretrained_tag(pretrained_model):
|
||||
pretrained_model = pretrained_model.lower()
|
||||
if "laion" in pretrained_model or "open_clip" in pretrained_model:
|
||||
return "open_clip"
|
||||
elif "openai" in pretrained_model:
|
||||
return "clip"
|
||||
elif "eva" in pretrained_model and "clip" in pretrained_model:
|
||||
return "eva_clip"
|
||||
else:
|
||||
return "other"
|
||||
|
||||
def load_pretrained_checkpoint(
|
||||
model,
|
||||
visual_checkpoint_path,
|
||||
text_checkpoint_path,
|
||||
strict=True,
|
||||
visual_model=None,
|
||||
text_model=None,
|
||||
model_key="model|module|state_dict",
|
||||
skip_list=[]):
|
||||
visual_tag = get_pretrained_tag(visual_model)
|
||||
text_tag = get_pretrained_tag(text_model)
|
||||
|
||||
logging.info(f"num of model state_dict keys: {len(model.state_dict().keys())}")
|
||||
visual_incompatible_keys, text_incompatible_keys = None, None
|
||||
if visual_checkpoint_path:
|
||||
if visual_tag == "eva_clip" or visual_tag == "open_clip":
|
||||
visual_state_dict = load_clip_visual_state_dict(visual_checkpoint_path, is_openai=False, skip_list=skip_list)
|
||||
elif visual_tag == "clip":
|
||||
visual_state_dict = load_clip_visual_state_dict(visual_checkpoint_path, is_openai=True, skip_list=skip_list)
|
||||
else:
|
||||
visual_state_dict = load_state_dict(visual_checkpoint_path, model_key=model_key, is_openai=False, skip_list=skip_list)
|
||||
|
||||
# resize_clip_pos_embed for CLIP and open CLIP
|
||||
if 'positional_embedding' in visual_state_dict:
|
||||
resize_visual_pos_embed(visual_state_dict, model)
|
||||
# specified to EVA model
|
||||
elif 'pos_embed' in visual_state_dict:
|
||||
resize_eva_pos_embed(visual_state_dict, model)
|
||||
|
||||
visual_incompatible_keys = model.visual.load_state_dict(visual_state_dict, strict=strict)
|
||||
logging.info(f"num of loaded visual_state_dict keys: {len(visual_state_dict.keys())}")
|
||||
logging.info(f"visual_incompatible_keys.missing_keys: {visual_incompatible_keys.missing_keys}")
|
||||
|
||||
if text_checkpoint_path:
|
||||
if text_tag == "eva_clip" or text_tag == "open_clip":
|
||||
text_state_dict = load_clip_text_state_dict(text_checkpoint_path, is_openai=False, skip_list=skip_list)
|
||||
elif text_tag == "clip":
|
||||
text_state_dict = load_clip_text_state_dict(text_checkpoint_path, is_openai=True, skip_list=skip_list)
|
||||
else:
|
||||
text_state_dict = load_state_dict(visual_checkpoint_path, model_key=model_key, is_openai=False, skip_list=skip_list)
|
||||
|
||||
text_incompatible_keys = model.text.load_state_dict(text_state_dict, strict=strict)
|
||||
|
||||
logging.info(f"num of loaded text_state_dict keys: {len(text_state_dict.keys())}")
|
||||
logging.info(f"text_incompatible_keys.missing_keys: {text_incompatible_keys.missing_keys}")
|
||||
|
||||
return visual_incompatible_keys, text_incompatible_keys
|
||||
|
||||
def create_model(
|
||||
model_name: str,
|
||||
pretrained: Optional[str] = None,
|
||||
precision: str = 'fp32',
|
||||
device: Union[str, torch.device] = 'cpu',
|
||||
jit: bool = False,
|
||||
force_quick_gelu: bool = False,
|
||||
force_custom_clip: bool = False,
|
||||
force_patch_dropout: Optional[float] = None,
|
||||
pretrained_image: str = '',
|
||||
pretrained_text: str = '',
|
||||
pretrained_hf: bool = True,
|
||||
pretrained_visual_model: str = None,
|
||||
pretrained_text_model: str = None,
|
||||
cache_dir: Optional[str] = None,
|
||||
skip_list: list = [],
|
||||
is_only_visual: bool = False,
|
||||
is_only_text: bool = False,
|
||||
):
|
||||
model_name = model_name.replace('/', '-') # for callers using old naming with / in ViT names
|
||||
if isinstance(device, str):
|
||||
device = torch.device(device)
|
||||
|
||||
if pretrained and pretrained.lower() == 'openai':
|
||||
logging.info(f'Loading pretrained {model_name} from OpenAI.')
|
||||
model = load_openai_model(
|
||||
model_name,
|
||||
precision=precision,
|
||||
device=device,
|
||||
jit=jit,
|
||||
cache_dir=cache_dir,
|
||||
)
|
||||
else:
|
||||
model_cfg = get_model_config(model_name)
|
||||
if model_cfg is not None:
|
||||
logging.info(f'Loaded {model_name} model config.')
|
||||
else:
|
||||
logging.error(f'Model config for {model_name} not found; available models {list_models()}.')
|
||||
raise RuntimeError(f'Model config for {model_name} not found.')
|
||||
|
||||
if 'rope' in model_cfg.get('vision_cfg', {}):
|
||||
if model_cfg['vision_cfg']['rope']:
|
||||
os.environ['RoPE'] = "1"
|
||||
else:
|
||||
os.environ['RoPE'] = "0"
|
||||
|
||||
if force_quick_gelu:
|
||||
# override for use of QuickGELU on non-OpenAI transformer models
|
||||
model_cfg["quick_gelu"] = True
|
||||
|
||||
if force_patch_dropout is not None:
|
||||
# override the default patch dropout value
|
||||
model_cfg['vision_cfg']["patch_dropout"] = force_patch_dropout
|
||||
|
||||
cast_dtype = get_cast_dtype(precision)
|
||||
custom_clip = model_cfg.pop('custom_text', False) or force_custom_clip or ('hf_model_name' in model_cfg['text_cfg'])
|
||||
|
||||
|
||||
if custom_clip:
|
||||
if 'hf_model_name' in model_cfg.get('text_cfg', {}):
|
||||
model_cfg['text_cfg']['hf_model_pretrained'] = pretrained_hf
|
||||
model = CustomCLIP(**model_cfg, cast_dtype=cast_dtype, is_only_visual=is_only_visual, is_only_text=is_only_text)
|
||||
else:
|
||||
model = CLIP(**model_cfg, cast_dtype=cast_dtype)
|
||||
print("Not CustomCLIP: If you have set building only visual or text tower, you may still get a complete CLIP model.")
|
||||
|
||||
pretrained_cfg = {}
|
||||
if pretrained:
|
||||
checkpoint_path = ''
|
||||
pretrained_cfg = get_pretrained_cfg(model_name, pretrained)
|
||||
if pretrained_cfg:
|
||||
checkpoint_path = download_pretrained(pretrained_cfg, cache_dir=cache_dir)
|
||||
elif os.path.exists(pretrained):
|
||||
checkpoint_path = pretrained
|
||||
|
||||
if checkpoint_path:
|
||||
logging.info(f'Loading pretrained {model_name} weights ({pretrained}).')
|
||||
load_checkpoint(model,
|
||||
checkpoint_path,
|
||||
model_key="model|module|state_dict",
|
||||
strict=False
|
||||
)
|
||||
else:
|
||||
error_str = (
|
||||
f'Pretrained weights ({pretrained}) not found for model {model_name}.'
|
||||
f'Available pretrained tags ({list_pretrained_tags_by_model(model_name)}.')
|
||||
logging.warning(error_str)
|
||||
raise RuntimeError(error_str)
|
||||
else:
|
||||
visual_checkpoint_path = ''
|
||||
text_checkpoint_path = ''
|
||||
|
||||
if pretrained_image:
|
||||
pretrained_visual_model = pretrained_visual_model.replace('/', '-') # for callers using old naming with / in ViT names
|
||||
pretrained_image_cfg = get_pretrained_cfg(pretrained_visual_model, pretrained_image)
|
||||
if 'timm_model_name' in model_cfg.get('vision_cfg', {}):
|
||||
# pretrained weight loading for timm models set via vision_cfg
|
||||
model_cfg['vision_cfg']['timm_model_pretrained'] = True
|
||||
elif pretrained_image_cfg:
|
||||
visual_checkpoint_path = download_pretrained(pretrained_image_cfg, cache_dir=cache_dir)
|
||||
elif os.path.exists(pretrained_image):
|
||||
visual_checkpoint_path = pretrained_image
|
||||
else:
|
||||
logging.warning(f'Pretrained weights ({visual_checkpoint_path}) not found for model {model_name}.visual.')
|
||||
raise RuntimeError(f'Pretrained weights ({visual_checkpoint_path}) not found for model {model_name}.visual.')
|
||||
|
||||
if pretrained_text:
|
||||
pretrained_text_model = pretrained_text_model.replace('/', '-') # for callers using old naming with / in ViT names
|
||||
pretrained_text_cfg = get_pretrained_cfg(pretrained_text_model, pretrained_text)
|
||||
if pretrained_image_cfg:
|
||||
text_checkpoint_path = download_pretrained(pretrained_text_cfg, cache_dir=cache_dir)
|
||||
elif os.path.exists(pretrained_text):
|
||||
text_checkpoint_path = pretrained_text
|
||||
else:
|
||||
logging.warning(f'Pretrained weights ({text_checkpoint_path}) not found for model {model_name}.text.')
|
||||
raise RuntimeError(f'Pretrained weights ({text_checkpoint_path}) not found for model {model_name}.text.')
|
||||
|
||||
if visual_checkpoint_path:
|
||||
logging.info(f'Loading pretrained {model_name}.visual weights ({visual_checkpoint_path}).')
|
||||
if text_checkpoint_path:
|
||||
logging.info(f'Loading pretrained {model_name}.text weights ({text_checkpoint_path}).')
|
||||
|
||||
if visual_checkpoint_path or text_checkpoint_path:
|
||||
load_pretrained_checkpoint(
|
||||
model,
|
||||
visual_checkpoint_path,
|
||||
text_checkpoint_path,
|
||||
strict=False,
|
||||
visual_model=pretrained_visual_model,
|
||||
text_model=pretrained_text_model,
|
||||
model_key="model|module|state_dict",
|
||||
skip_list=skip_list
|
||||
)
|
||||
|
||||
if "fp16" in precision or "bf16" in precision:
|
||||
logging.info(f'convert precision to {precision}')
|
||||
model = model.to(torch.bfloat16) if 'bf16' in precision else model.to(torch.float16)
|
||||
|
||||
model.to(device=device)
|
||||
|
||||
# set image / mean metadata from pretrained_cfg if available, or use default
|
||||
if not is_only_text:
|
||||
model.visual.image_mean = pretrained_cfg.get('mean', None) or OPENAI_DATASET_MEAN
|
||||
model.visual.image_std = pretrained_cfg.get('std', None) or OPENAI_DATASET_STD
|
||||
|
||||
if jit:
|
||||
model = torch.jit.script(model)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def create_model_and_transforms(
|
||||
model_name: str,
|
||||
pretrained: Optional[str] = None,
|
||||
precision: str = 'fp32',
|
||||
device: Union[str, torch.device] = 'cpu',
|
||||
jit: bool = False,
|
||||
force_quick_gelu: bool = False,
|
||||
force_custom_clip: bool = False,
|
||||
force_patch_dropout: Optional[float] = None,
|
||||
pretrained_image: str = '',
|
||||
pretrained_text: str = '',
|
||||
pretrained_hf: bool = True,
|
||||
pretrained_visual_model: str = None,
|
||||
pretrained_text_model: str = None,
|
||||
image_mean: Optional[Tuple[float, ...]] = None,
|
||||
image_std: Optional[Tuple[float, ...]] = None,
|
||||
cache_dir: Optional[str] = None,
|
||||
skip_list: list = [],
|
||||
):
|
||||
model = create_model(
|
||||
model_name,
|
||||
pretrained,
|
||||
precision=precision,
|
||||
device=device,
|
||||
jit=jit,
|
||||
force_quick_gelu=force_quick_gelu,
|
||||
force_custom_clip=force_custom_clip,
|
||||
force_patch_dropout=force_patch_dropout,
|
||||
pretrained_image=pretrained_image,
|
||||
pretrained_text=pretrained_text,
|
||||
pretrained_hf=pretrained_hf,
|
||||
pretrained_visual_model=pretrained_visual_model,
|
||||
pretrained_text_model=pretrained_text_model,
|
||||
cache_dir=cache_dir,
|
||||
skip_list=skip_list,
|
||||
)
|
||||
|
||||
image_mean = image_mean or getattr(model.visual, 'image_mean', None)
|
||||
image_std = image_std or getattr(model.visual, 'image_std', None)
|
||||
preprocess_train = image_transform(
|
||||
model.visual.image_size,
|
||||
is_train=True,
|
||||
mean=image_mean,
|
||||
std=image_std
|
||||
)
|
||||
preprocess_val = image_transform(
|
||||
model.visual.image_size,
|
||||
is_train=False,
|
||||
mean=image_mean,
|
||||
std=image_std
|
||||
)
|
||||
|
||||
return model, preprocess_train, preprocess_val
|
||||
|
||||
def create_eva_vision_and_transforms(
|
||||
model_name: str,
|
||||
pretrained: Optional[str] = None,
|
||||
precision: str = 'fp32',
|
||||
device: Union[str, torch.device] = 'cpu',
|
||||
jit: bool = False,
|
||||
force_quick_gelu: bool = False,
|
||||
force_custom_clip: bool = False,
|
||||
force_patch_dropout: Optional[float] = None,
|
||||
pretrained_image: str = '',
|
||||
pretrained_text: str = '',
|
||||
pretrained_hf: bool = True,
|
||||
pretrained_visual_model: str = None,
|
||||
pretrained_text_model: str = None,
|
||||
image_mean: Optional[Tuple[float, ...]] = None,
|
||||
image_std: Optional[Tuple[float, ...]] = None,
|
||||
cache_dir: Optional[str] = None,
|
||||
skip_list: list = [],
|
||||
):
|
||||
model = create_model(
|
||||
model_name,
|
||||
pretrained,
|
||||
precision=precision,
|
||||
device=device,
|
||||
jit=jit,
|
||||
force_quick_gelu=force_quick_gelu,
|
||||
force_custom_clip=force_custom_clip,
|
||||
force_patch_dropout=force_patch_dropout,
|
||||
pretrained_image=pretrained_image,
|
||||
pretrained_text=pretrained_text,
|
||||
pretrained_hf=pretrained_hf,
|
||||
pretrained_visual_model=pretrained_visual_model,
|
||||
pretrained_text_model=pretrained_text_model,
|
||||
cache_dir=cache_dir,
|
||||
skip_list=skip_list,
|
||||
is_only_visual=True, # only use visual tower
|
||||
)
|
||||
|
||||
image_mean = image_mean or getattr(model.visual, 'image_mean', None)
|
||||
image_std = image_std or getattr(model.visual, 'image_std', None)
|
||||
preprocess_train = image_transform(
|
||||
model.visual.image_size,
|
||||
is_train=True,
|
||||
mean=image_mean,
|
||||
std=image_std
|
||||
)
|
||||
preprocess_val = image_transform(
|
||||
model.visual.image_size,
|
||||
is_train=False,
|
||||
mean=image_mean,
|
||||
std=image_std
|
||||
)
|
||||
|
||||
return model, preprocess_train, preprocess_val
|
||||
|
||||
def create_model_from_pretrained(
|
||||
model_name: str,
|
||||
pretrained: str,
|
||||
precision: str = 'fp32',
|
||||
device: Union[str, torch.device] = 'cpu',
|
||||
jit: bool = False,
|
||||
force_quick_gelu: bool = False,
|
||||
force_custom_clip: bool = False,
|
||||
force_patch_dropout: Optional[float] = None,
|
||||
return_transform: bool = True,
|
||||
image_mean: Optional[Tuple[float, ...]] = None,
|
||||
image_std: Optional[Tuple[float, ...]] = None,
|
||||
cache_dir: Optional[str] = None,
|
||||
is_frozen: bool = False,
|
||||
):
|
||||
if not is_pretrained_cfg(model_name, pretrained) and not os.path.exists(pretrained):
|
||||
raise RuntimeError(
|
||||
f'{pretrained} is not a valid pretrained cfg or checkpoint for {model_name}.'
|
||||
f' Use open_clip.list_pretrained() to find one.')
|
||||
|
||||
model = create_model(
|
||||
model_name,
|
||||
pretrained,
|
||||
precision=precision,
|
||||
device=device,
|
||||
jit=jit,
|
||||
force_quick_gelu=force_quick_gelu,
|
||||
force_custom_clip=force_custom_clip,
|
||||
force_patch_dropout=force_patch_dropout,
|
||||
cache_dir=cache_dir,
|
||||
)
|
||||
|
||||
if is_frozen:
|
||||
for param in model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
if not return_transform:
|
||||
return model
|
||||
|
||||
image_mean = image_mean or getattr(model.visual, 'image_mean', None)
|
||||
image_std = image_std or getattr(model.visual, 'image_std', None)
|
||||
preprocess = image_transform(
|
||||
model.visual.image_size,
|
||||
is_train=False,
|
||||
mean=image_mean,
|
||||
std=image_std
|
||||
)
|
||||
|
||||
return model, preprocess
|
||||
@@ -0,0 +1,57 @@
|
||||
# HF architecture dict:
|
||||
arch_dict = {
|
||||
# https://huggingface.co/docs/transformers/model_doc/roberta#roberta
|
||||
"roberta": {
|
||||
"config_names": {
|
||||
"context_length": "max_position_embeddings",
|
||||
"vocab_size": "vocab_size",
|
||||
"width": "hidden_size",
|
||||
"heads": "num_attention_heads",
|
||||
"layers": "num_hidden_layers",
|
||||
"layer_attr": "layer",
|
||||
"token_embeddings_attr": "embeddings"
|
||||
},
|
||||
"pooler": "mean_pooler",
|
||||
},
|
||||
# https://huggingface.co/docs/transformers/model_doc/xlm-roberta#transformers.XLMRobertaConfig
|
||||
"xlm-roberta": {
|
||||
"config_names": {
|
||||
"context_length": "max_position_embeddings",
|
||||
"vocab_size": "vocab_size",
|
||||
"width": "hidden_size",
|
||||
"heads": "num_attention_heads",
|
||||
"layers": "num_hidden_layers",
|
||||
"layer_attr": "layer",
|
||||
"token_embeddings_attr": "embeddings"
|
||||
},
|
||||
"pooler": "mean_pooler",
|
||||
},
|
||||
# https://huggingface.co/docs/transformers/model_doc/mt5#mt5
|
||||
"mt5": {
|
||||
"config_names": {
|
||||
# unlimited seqlen
|
||||
# https://github.com/google-research/text-to-text-transfer-transformer/issues/273
|
||||
# https://github.com/huggingface/transformers/blob/v4.24.0/src/transformers/models/t5/modeling_t5.py#L374
|
||||
"context_length": "",
|
||||
"vocab_size": "vocab_size",
|
||||
"width": "d_model",
|
||||
"heads": "num_heads",
|
||||
"layers": "num_layers",
|
||||
"layer_attr": "block",
|
||||
"token_embeddings_attr": "embed_tokens"
|
||||
},
|
||||
"pooler": "mean_pooler",
|
||||
},
|
||||
"bert": {
|
||||
"config_names": {
|
||||
"context_length": "max_position_embeddings",
|
||||
"vocab_size": "vocab_size",
|
||||
"width": "hidden_size",
|
||||
"heads": "num_attention_heads",
|
||||
"layers": "num_hidden_layers",
|
||||
"layer_attr": "layer",
|
||||
"token_embeddings_attr": "embeddings"
|
||||
},
|
||||
"pooler": "mean_pooler",
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,248 @@
|
||||
""" huggingface model adapter
|
||||
|
||||
Wraps HuggingFace transformers (https://github.com/huggingface/transformers) models for use as a text tower in CLIP model.
|
||||
"""
|
||||
|
||||
import re
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
from torch import TensorType
|
||||
try:
|
||||
import transformers
|
||||
from transformers import AutoModel, AutoModelForMaskedLM, AutoTokenizer, AutoConfig, PretrainedConfig
|
||||
from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, \
|
||||
BaseModelOutputWithPoolingAndCrossAttentions
|
||||
except ImportError as e:
|
||||
transformers = None
|
||||
|
||||
|
||||
class BaseModelOutput:
|
||||
pass
|
||||
|
||||
|
||||
class PretrainedConfig:
|
||||
pass
|
||||
|
||||
from .hf_configs import arch_dict
|
||||
|
||||
# utils
|
||||
def _camel2snake(s):
|
||||
return re.sub(r'(?<!^)(?=[A-Z])', '_', s).lower()
|
||||
|
||||
# TODO: ?last - for gpt-like models
|
||||
_POOLERS = {}
|
||||
|
||||
def register_pooler(cls):
|
||||
"""Decorator registering pooler class"""
|
||||
_POOLERS[_camel2snake(cls.__name__)] = cls
|
||||
return cls
|
||||
|
||||
|
||||
@register_pooler
|
||||
class MeanPooler(nn.Module):
|
||||
"""Mean pooling"""
|
||||
def forward(self, x:BaseModelOutput, attention_mask:TensorType):
|
||||
masked_output = x.last_hidden_state * attention_mask.unsqueeze(-1)
|
||||
return masked_output.sum(dim=1) / attention_mask.sum(-1, keepdim=True)
|
||||
|
||||
@register_pooler
|
||||
class MaxPooler(nn.Module):
|
||||
"""Max pooling"""
|
||||
def forward(self, x:BaseModelOutput, attention_mask:TensorType):
|
||||
masked_output = x.last_hidden_state.masked_fill(attention_mask.unsqueeze(-1), -torch.inf)
|
||||
return masked_output.max(1).values
|
||||
|
||||
@register_pooler
|
||||
class ClsPooler(nn.Module):
|
||||
"""CLS token pooling"""
|
||||
def __init__(self, use_pooler_output=True):
|
||||
super().__init__()
|
||||
self.cls_token_position = 0
|
||||
self.use_pooler_output = use_pooler_output
|
||||
|
||||
def forward(self, x:BaseModelOutput, attention_mask:TensorType):
|
||||
|
||||
if (self.use_pooler_output and
|
||||
isinstance(x, (BaseModelOutputWithPooling, BaseModelOutputWithPoolingAndCrossAttentions)) and
|
||||
(x.pooler_output is not None)
|
||||
):
|
||||
return x.pooler_output
|
||||
|
||||
return x.last_hidden_state[:, self.cls_token_position, :]
|
||||
|
||||
class HFTextEncoder(nn.Module):
|
||||
"""HuggingFace model adapter"""
|
||||
def __init__(
|
||||
self,
|
||||
model_name_or_path: str,
|
||||
output_dim: int,
|
||||
tokenizer_name: str = None,
|
||||
config: PretrainedConfig = None,
|
||||
pooler_type: str = None,
|
||||
proj: str = None,
|
||||
pretrained: bool = True,
|
||||
masked_language_modeling: bool = False):
|
||||
super().__init__()
|
||||
|
||||
self.output_dim = output_dim
|
||||
|
||||
# TODO: find better way to get this information
|
||||
uses_transformer_pooler = (pooler_type == "cls_pooler")
|
||||
|
||||
if transformers is None:
|
||||
raise RuntimeError("Please `pip install transformers` to use pre-trained HuggingFace models")
|
||||
if config is None:
|
||||
self.config = AutoConfig.from_pretrained(model_name_or_path)
|
||||
if masked_language_modeling:
|
||||
create_func, model_args = (AutoModelForMaskedLM.from_pretrained, model_name_or_path) if pretrained else (
|
||||
AutoModelForMaskedLM.from_config, self.config)
|
||||
else:
|
||||
create_func, model_args = (AutoModel.from_pretrained, model_name_or_path) if pretrained else (
|
||||
AutoModel.from_config, self.config)
|
||||
# TODO: do all model configs have this attribute? PretrainedConfig does so yes??
|
||||
if hasattr(self.config, "is_encoder_decoder") and self.config.is_encoder_decoder:
|
||||
self.transformer = create_func(model_args)
|
||||
self.transformer = self.transformer.encoder
|
||||
else:
|
||||
self.transformer = create_func(model_args, add_pooling_layer=uses_transformer_pooler)
|
||||
else:
|
||||
self.config = config
|
||||
if masked_language_modeling:
|
||||
self.transformer = AutoModelForMaskedLM.from_config(config)
|
||||
else:
|
||||
self.transformer = AutoModel.from_config(config)
|
||||
|
||||
if pooler_type is None: # get default arch pooler
|
||||
self.pooler = _POOLERS[(arch_dict[self.config.model_type]["pooler"])]()
|
||||
else:
|
||||
self.pooler = _POOLERS[pooler_type]()
|
||||
|
||||
d_model = getattr(self.config, arch_dict[self.config.model_type]["config_names"]["width"])
|
||||
if (d_model == output_dim) and (proj is None): # do we always need a proj?
|
||||
self.proj = nn.Identity()
|
||||
elif proj == 'linear':
|
||||
self.proj = nn.Linear(d_model, output_dim, bias=False)
|
||||
elif proj == 'mlp':
|
||||
hidden_size = (d_model + output_dim) // 2
|
||||
self.proj = nn.Sequential(
|
||||
nn.Linear(d_model, hidden_size, bias=False),
|
||||
nn.GELU(),
|
||||
nn.Linear(hidden_size, output_dim, bias=False),
|
||||
)
|
||||
|
||||
# self.itm_proj = nn.Linear(d_model, 2, bias=False)
|
||||
# self.mlm_proj = nn.Linear(d_model, self.config.vocab_size), bias=False)
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
|
||||
|
||||
# def forward_itm(self, x:TensorType, image_embeds:TensorType) -> TensorType:
|
||||
# image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(x.device)
|
||||
# attn_mask = (x != self.config.pad_token_id).long()
|
||||
# out = self.transformer(
|
||||
# input_ids=x,
|
||||
# attention_mask=attn_mask,
|
||||
# encoder_hidden_states = image_embeds,
|
||||
# encoder_attention_mask = image_atts,
|
||||
# )
|
||||
# pooled_out = self.pooler(out, attn_mask)
|
||||
|
||||
# return self.itm_proj(pooled_out)
|
||||
|
||||
def mask(self, input_ids, vocab_size, device, targets=None, masked_indices=None, probability_matrix=None):
|
||||
if masked_indices is None:
|
||||
masked_indices = torch.bernoulli(probability_matrix).bool()
|
||||
|
||||
masked_indices[input_ids == self.tokenizer.pad_token_id] = False
|
||||
masked_indices[input_ids == self.tokenizer.cls_token_id] = False
|
||||
|
||||
if targets is not None:
|
||||
targets[~masked_indices] = -100 # We only compute loss on masked tokens
|
||||
|
||||
# 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
|
||||
indices_replaced = torch.bernoulli(torch.full(input_ids.shape, 0.8)).bool() & masked_indices
|
||||
input_ids[indices_replaced] = self.tokenizer.mask_token_id
|
||||
|
||||
# 10% of the time, we replace masked input tokens with random word
|
||||
indices_random = torch.bernoulli(torch.full(input_ids.shape, 0.5)).bool() & masked_indices & ~indices_replaced
|
||||
random_words = torch.randint(vocab_size, input_ids.shape, dtype=torch.long).to(device)
|
||||
input_ids[indices_random] = random_words[indices_random]
|
||||
# The rest of the time (10% of the time) we keep the masked input tokens unchanged
|
||||
|
||||
if targets is not None:
|
||||
return input_ids, targets
|
||||
else:
|
||||
return input_ids
|
||||
|
||||
def forward_mlm(self, input_ids, image_embeds, mlm_probability=0.25):
|
||||
labels = input_ids.clone()
|
||||
attn_mask = (input_ids != self.config.pad_token_id).long()
|
||||
image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(input_ids.device)
|
||||
vocab_size = getattr(self.config, arch_dict[self.config.model_type]["config_names"]["vocab_size"])
|
||||
probability_matrix = torch.full(labels.shape, mlm_probability)
|
||||
input_ids, labels = self.mask(input_ids, vocab_size, input_ids.device, targets=labels,
|
||||
probability_matrix = probability_matrix)
|
||||
mlm_output = self.transformer(input_ids,
|
||||
attention_mask = attn_mask,
|
||||
encoder_hidden_states = image_embeds,
|
||||
encoder_attention_mask = image_atts,
|
||||
return_dict = True,
|
||||
labels = labels,
|
||||
)
|
||||
return mlm_output.loss
|
||||
# mlm_output = self.transformer(input_ids,
|
||||
# attention_mask = attn_mask,
|
||||
# encoder_hidden_states = image_embeds,
|
||||
# encoder_attention_mask = image_atts,
|
||||
# return_dict = True,
|
||||
# ).last_hidden_state
|
||||
# logits = self.mlm_proj(mlm_output)
|
||||
|
||||
# # logits = logits[:, :-1, :].contiguous().view(-1, vocab_size)
|
||||
# logits = logits[:, 1:, :].contiguous().view(-1, vocab_size)
|
||||
# labels = labels[:, 1:].contiguous().view(-1)
|
||||
|
||||
# mlm_loss = F.cross_entropy(
|
||||
# logits,
|
||||
# labels,
|
||||
# # label_smoothing=0.1,
|
||||
# )
|
||||
# return mlm_loss
|
||||
|
||||
|
||||
def forward(self, x:TensorType) -> TensorType:
|
||||
attn_mask = (x != self.config.pad_token_id).long()
|
||||
out = self.transformer(input_ids=x, attention_mask=attn_mask)
|
||||
pooled_out = self.pooler(out, attn_mask)
|
||||
|
||||
return self.proj(pooled_out)
|
||||
|
||||
def lock(self, unlocked_layers:int=0, freeze_layer_norm:bool=True):
|
||||
if not unlocked_layers: # full freezing
|
||||
for n, p in self.transformer.named_parameters():
|
||||
p.requires_grad = (not freeze_layer_norm) if "LayerNorm" in n.split(".") else False
|
||||
return
|
||||
|
||||
encoder = self.transformer.encoder if hasattr(self.transformer, 'encoder') else self.transformer
|
||||
layer_list = getattr(encoder, arch_dict[self.config.model_type]["config_names"]["layer_attr"])
|
||||
print(f"Unlocking {unlocked_layers}/{len(layer_list) + 1} layers of hf model")
|
||||
embeddings = getattr(
|
||||
self.transformer, arch_dict[self.config.model_type]["config_names"]["token_embeddings_attr"])
|
||||
modules = [embeddings, *layer_list][:-unlocked_layers]
|
||||
# freeze layers
|
||||
for module in modules:
|
||||
for n, p in module.named_parameters():
|
||||
p.requires_grad = (not freeze_layer_norm) if "LayerNorm" in n.split(".") else False
|
||||
|
||||
|
||||
@torch.jit.ignore
|
||||
def set_grad_checkpointing(self, enable=True):
|
||||
self.transformer.gradient_checkpointing_enable()
|
||||
|
||||
def get_num_layers(self):
|
||||
encoder = self.transformer.encoder if hasattr(self.transformer, 'encoder') else self.transformer
|
||||
layer_list = getattr(encoder, arch_dict[self.config.model_type]["config_names"]["layer_attr"])
|
||||
return len(layer_list)
|
||||
|
||||
def init_parameters(self):
|
||||
pass
|
||||
@@ -0,0 +1,138 @@
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
try:
|
||||
import torch.distributed.nn
|
||||
from torch import distributed as dist
|
||||
has_distributed = True
|
||||
except ImportError:
|
||||
has_distributed = False
|
||||
|
||||
try:
|
||||
import horovod.torch as hvd
|
||||
except ImportError:
|
||||
hvd = None
|
||||
|
||||
from timm.loss import LabelSmoothingCrossEntropy
|
||||
|
||||
|
||||
def gather_features(
|
||||
image_features,
|
||||
text_features,
|
||||
local_loss=False,
|
||||
gather_with_grad=False,
|
||||
rank=0,
|
||||
world_size=1,
|
||||
use_horovod=False
|
||||
):
|
||||
assert has_distributed, 'torch.distributed did not import correctly, please use a PyTorch version with support.'
|
||||
if use_horovod:
|
||||
assert hvd is not None, 'Please install horovod'
|
||||
if gather_with_grad:
|
||||
all_image_features = hvd.allgather(image_features)
|
||||
all_text_features = hvd.allgather(text_features)
|
||||
else:
|
||||
with torch.no_grad():
|
||||
all_image_features = hvd.allgather(image_features)
|
||||
all_text_features = hvd.allgather(text_features)
|
||||
if not local_loss:
|
||||
# ensure grads for local rank when all_* features don't have a gradient
|
||||
gathered_image_features = list(all_image_features.chunk(world_size, dim=0))
|
||||
gathered_text_features = list(all_text_features.chunk(world_size, dim=0))
|
||||
gathered_image_features[rank] = image_features
|
||||
gathered_text_features[rank] = text_features
|
||||
all_image_features = torch.cat(gathered_image_features, dim=0)
|
||||
all_text_features = torch.cat(gathered_text_features, dim=0)
|
||||
else:
|
||||
# We gather tensors from all gpus
|
||||
if gather_with_grad:
|
||||
all_image_features = torch.cat(torch.distributed.nn.all_gather(image_features), dim=0)
|
||||
all_text_features = torch.cat(torch.distributed.nn.all_gather(text_features), dim=0)
|
||||
# all_image_features = torch.cat(torch.distributed.nn.all_gather(image_features, async_op=True), dim=0)
|
||||
# all_text_features = torch.cat(torch.distributed.nn.all_gather(text_features, async_op=True), dim=0)
|
||||
else:
|
||||
gathered_image_features = [torch.zeros_like(image_features) for _ in range(world_size)]
|
||||
gathered_text_features = [torch.zeros_like(text_features) for _ in range(world_size)]
|
||||
dist.all_gather(gathered_image_features, image_features)
|
||||
dist.all_gather(gathered_text_features, text_features)
|
||||
if not local_loss:
|
||||
# ensure grads for local rank when all_* features don't have a gradient
|
||||
gathered_image_features[rank] = image_features
|
||||
gathered_text_features[rank] = text_features
|
||||
all_image_features = torch.cat(gathered_image_features, dim=0)
|
||||
all_text_features = torch.cat(gathered_text_features, dim=0)
|
||||
|
||||
return all_image_features, all_text_features
|
||||
|
||||
|
||||
class ClipLoss(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
local_loss=False,
|
||||
gather_with_grad=False,
|
||||
cache_labels=False,
|
||||
rank=0,
|
||||
world_size=1,
|
||||
use_horovod=False,
|
||||
smoothing=0.,
|
||||
):
|
||||
super().__init__()
|
||||
self.local_loss = local_loss
|
||||
self.gather_with_grad = gather_with_grad
|
||||
self.cache_labels = cache_labels
|
||||
self.rank = rank
|
||||
self.world_size = world_size
|
||||
self.use_horovod = use_horovod
|
||||
self.label_smoothing_cross_entropy = LabelSmoothingCrossEntropy(smoothing=smoothing) if smoothing > 0 else None
|
||||
|
||||
# cache state
|
||||
self.prev_num_logits = 0
|
||||
self.labels = {}
|
||||
|
||||
def forward(self, image_features, text_features, logit_scale=1.):
|
||||
device = image_features.device
|
||||
if self.world_size > 1:
|
||||
all_image_features, all_text_features = gather_features(
|
||||
image_features, text_features,
|
||||
self.local_loss, self.gather_with_grad, self.rank, self.world_size, self.use_horovod)
|
||||
|
||||
if self.local_loss:
|
||||
logits_per_image = logit_scale * image_features @ all_text_features.T
|
||||
logits_per_text = logit_scale * text_features @ all_image_features.T
|
||||
else:
|
||||
logits_per_image = logit_scale * all_image_features @ all_text_features.T
|
||||
logits_per_text = logits_per_image.T
|
||||
else:
|
||||
logits_per_image = logit_scale * image_features @ text_features.T
|
||||
logits_per_text = logit_scale * text_features @ image_features.T
|
||||
# calculated ground-truth and cache if enabled
|
||||
num_logits = logits_per_image.shape[0]
|
||||
if self.prev_num_logits != num_logits or device not in self.labels:
|
||||
labels = torch.arange(num_logits, device=device, dtype=torch.long)
|
||||
if self.world_size > 1 and self.local_loss:
|
||||
labels = labels + num_logits * self.rank
|
||||
if self.cache_labels:
|
||||
self.labels[device] = labels
|
||||
self.prev_num_logits = num_logits
|
||||
else:
|
||||
labels = self.labels[device]
|
||||
|
||||
if self.label_smoothing_cross_entropy:
|
||||
total_loss = (
|
||||
self.label_smoothing_cross_entropy(logits_per_image, labels) +
|
||||
self.label_smoothing_cross_entropy(logits_per_text, labels)
|
||||
) / 2
|
||||
else:
|
||||
total_loss = (
|
||||
F.cross_entropy(logits_per_image, labels) +
|
||||
F.cross_entropy(logits_per_text, labels)
|
||||
) / 2
|
||||
|
||||
acc = None
|
||||
i2t_acc = (logits_per_image.argmax(-1) == labels).sum() / len(logits_per_image)
|
||||
t2i_acc = (logits_per_text.argmax(-1) == labels).sum() / len(logits_per_text)
|
||||
acc = {"i2t": i2t_acc, "t2i": t2i_acc}
|
||||
return total_loss, acc
|
||||
@@ -0,0 +1,452 @@
|
||||
""" CLIP Model
|
||||
|
||||
Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
|
||||
"""
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Tuple, Union
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
try:
|
||||
from .hf_model import HFTextEncoder
|
||||
except:
|
||||
HFTextEncoder = None
|
||||
from .modified_resnet import ModifiedResNet
|
||||
from .timm_model import TimmModel
|
||||
from .eva_vit_model import EVAVisionTransformer
|
||||
from .transformer import LayerNorm, QuickGELU, Attention, VisionTransformer, TextTransformer
|
||||
|
||||
# try:
|
||||
# from apex.normalization import FusedLayerNorm
|
||||
# except:
|
||||
FusedLayerNorm = LayerNorm
|
||||
# print("Please 'pip install apex'")
|
||||
|
||||
try:
|
||||
import xformers.ops as xops
|
||||
except ImportError:
|
||||
xops = None
|
||||
# print("Please 'pip install xformers'")
|
||||
|
||||
@dataclass
|
||||
class CLIPVisionCfg:
|
||||
layers: Union[Tuple[int, int, int, int], int] = 12
|
||||
width: int = 768
|
||||
head_width: int = 64
|
||||
mlp_ratio: float = 4.0
|
||||
patch_size: int = 16
|
||||
image_size: Union[Tuple[int, int], int] = 224
|
||||
ls_init_value: Optional[float] = None # layer scale initial value
|
||||
patch_dropout: float = 0. # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
|
||||
global_average_pool: bool = False # whether to global average pool the last embedding layer, instead of using CLS token (https://arxiv.org/abs/2205.01580)
|
||||
drop_path_rate: Optional[float] = None # drop path rate
|
||||
timm_model_name: str = None # a valid model name overrides layers, width, patch_size
|
||||
timm_model_pretrained: bool = False # use (imagenet) pretrained weights for named model
|
||||
timm_pool: str = 'avg' # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
|
||||
timm_proj: str = 'linear' # linear projection for timm model output ('linear', 'mlp', '')
|
||||
timm_proj_bias: bool = False # enable bias final projection
|
||||
eva_model_name: str = None # a valid eva model name overrides layers, width, patch_size
|
||||
qkv_bias: bool = True
|
||||
fusedLN: bool = False
|
||||
xattn: bool = False
|
||||
postnorm: bool = False
|
||||
rope: bool = False
|
||||
pt_hw_seq_len: int = 16 # 224/14
|
||||
intp_freq: bool = False
|
||||
naiveswiglu: bool = False
|
||||
subln: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class CLIPTextCfg:
|
||||
context_length: int = 77
|
||||
vocab_size: int = 49408
|
||||
width: int = 512
|
||||
heads: int = 8
|
||||
layers: int = 12
|
||||
ls_init_value: Optional[float] = None # layer scale initial value
|
||||
hf_model_name: str = None
|
||||
hf_tokenizer_name: str = None
|
||||
hf_model_pretrained: bool = True
|
||||
proj: str = 'mlp'
|
||||
pooler_type: str = 'mean_pooler'
|
||||
masked_language_modeling: bool = False
|
||||
fusedLN: bool = False
|
||||
xattn: bool = False
|
||||
attn_mask: bool = True
|
||||
|
||||
def get_cast_dtype(precision: str):
|
||||
cast_dtype = None
|
||||
if precision == 'bf16':
|
||||
cast_dtype = torch.bfloat16
|
||||
elif precision == 'fp16':
|
||||
cast_dtype = torch.float16
|
||||
return cast_dtype
|
||||
|
||||
|
||||
def _build_vision_tower(
|
||||
embed_dim: int,
|
||||
vision_cfg: CLIPVisionCfg,
|
||||
quick_gelu: bool = False,
|
||||
cast_dtype: Optional[torch.dtype] = None
|
||||
):
|
||||
if isinstance(vision_cfg, dict):
|
||||
vision_cfg = CLIPVisionCfg(**vision_cfg)
|
||||
|
||||
# OpenAI models are pretrained w/ QuickGELU but native nn.GELU is both faster and more
|
||||
# memory efficient in recent PyTorch releases (>= 1.10).
|
||||
# NOTE: timm models always use native GELU regardless of quick_gelu flag.
|
||||
act_layer = QuickGELU if quick_gelu else nn.GELU
|
||||
|
||||
if vision_cfg.eva_model_name:
|
||||
vision_heads = vision_cfg.width // vision_cfg.head_width
|
||||
norm_layer = LayerNorm
|
||||
|
||||
visual = EVAVisionTransformer(
|
||||
img_size=vision_cfg.image_size,
|
||||
patch_size=vision_cfg.patch_size,
|
||||
num_classes=embed_dim,
|
||||
use_mean_pooling=vision_cfg.global_average_pool, #False
|
||||
init_values=vision_cfg.ls_init_value,
|
||||
patch_dropout=vision_cfg.patch_dropout,
|
||||
embed_dim=vision_cfg.width,
|
||||
depth=vision_cfg.layers,
|
||||
num_heads=vision_heads,
|
||||
mlp_ratio=vision_cfg.mlp_ratio,
|
||||
qkv_bias=vision_cfg.qkv_bias,
|
||||
drop_path_rate=vision_cfg.drop_path_rate,
|
||||
norm_layer= partial(FusedLayerNorm, eps=1e-6) if vision_cfg.fusedLN else partial(norm_layer, eps=1e-6),
|
||||
xattn=vision_cfg.xattn,
|
||||
rope=vision_cfg.rope,
|
||||
postnorm=vision_cfg.postnorm,
|
||||
pt_hw_seq_len= vision_cfg.pt_hw_seq_len, # 224/14
|
||||
intp_freq= vision_cfg.intp_freq,
|
||||
naiveswiglu= vision_cfg.naiveswiglu,
|
||||
subln= vision_cfg.subln
|
||||
)
|
||||
elif vision_cfg.timm_model_name:
|
||||
visual = TimmModel(
|
||||
vision_cfg.timm_model_name,
|
||||
pretrained=vision_cfg.timm_model_pretrained,
|
||||
pool=vision_cfg.timm_pool,
|
||||
proj=vision_cfg.timm_proj,
|
||||
proj_bias=vision_cfg.timm_proj_bias,
|
||||
embed_dim=embed_dim,
|
||||
image_size=vision_cfg.image_size
|
||||
)
|
||||
act_layer = nn.GELU # so that text transformer doesn't use QuickGELU w/ timm models
|
||||
elif isinstance(vision_cfg.layers, (tuple, list)):
|
||||
vision_heads = vision_cfg.width * 32 // vision_cfg.head_width
|
||||
visual = ModifiedResNet(
|
||||
layers=vision_cfg.layers,
|
||||
output_dim=embed_dim,
|
||||
heads=vision_heads,
|
||||
image_size=vision_cfg.image_size,
|
||||
width=vision_cfg.width
|
||||
)
|
||||
else:
|
||||
vision_heads = vision_cfg.width // vision_cfg.head_width
|
||||
norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
|
||||
visual = VisionTransformer(
|
||||
image_size=vision_cfg.image_size,
|
||||
patch_size=vision_cfg.patch_size,
|
||||
width=vision_cfg.width,
|
||||
layers=vision_cfg.layers,
|
||||
heads=vision_heads,
|
||||
mlp_ratio=vision_cfg.mlp_ratio,
|
||||
ls_init_value=vision_cfg.ls_init_value,
|
||||
patch_dropout=vision_cfg.patch_dropout,
|
||||
global_average_pool=vision_cfg.global_average_pool,
|
||||
output_dim=embed_dim,
|
||||
act_layer=act_layer,
|
||||
norm_layer=norm_layer,
|
||||
)
|
||||
|
||||
return visual
|
||||
|
||||
|
||||
def _build_text_tower(
|
||||
embed_dim: int,
|
||||
text_cfg: CLIPTextCfg,
|
||||
quick_gelu: bool = False,
|
||||
cast_dtype: Optional[torch.dtype] = None,
|
||||
):
|
||||
if isinstance(text_cfg, dict):
|
||||
text_cfg = CLIPTextCfg(**text_cfg)
|
||||
|
||||
if text_cfg.hf_model_name:
|
||||
text = HFTextEncoder(
|
||||
text_cfg.hf_model_name,
|
||||
output_dim=embed_dim,
|
||||
tokenizer_name=text_cfg.hf_tokenizer_name,
|
||||
proj=text_cfg.proj,
|
||||
pooler_type=text_cfg.pooler_type,
|
||||
masked_language_modeling=text_cfg.masked_language_modeling
|
||||
)
|
||||
else:
|
||||
act_layer = QuickGELU if quick_gelu else nn.GELU
|
||||
norm_layer = LayerNorm
|
||||
|
||||
text = TextTransformer(
|
||||
context_length=text_cfg.context_length,
|
||||
vocab_size=text_cfg.vocab_size,
|
||||
width=text_cfg.width,
|
||||
heads=text_cfg.heads,
|
||||
layers=text_cfg.layers,
|
||||
ls_init_value=text_cfg.ls_init_value,
|
||||
output_dim=embed_dim,
|
||||
act_layer=act_layer,
|
||||
norm_layer= FusedLayerNorm if text_cfg.fusedLN else norm_layer,
|
||||
xattn=text_cfg.xattn,
|
||||
attn_mask=text_cfg.attn_mask,
|
||||
)
|
||||
return text
|
||||
|
||||
class CLIP(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
embed_dim: int,
|
||||
vision_cfg: CLIPVisionCfg,
|
||||
text_cfg: CLIPTextCfg,
|
||||
quick_gelu: bool = False,
|
||||
cast_dtype: Optional[torch.dtype] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
|
||||
|
||||
text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
|
||||
self.transformer = text.transformer
|
||||
self.vocab_size = text.vocab_size
|
||||
self.token_embedding = text.token_embedding
|
||||
self.positional_embedding = text.positional_embedding
|
||||
self.ln_final = text.ln_final
|
||||
self.text_projection = text.text_projection
|
||||
self.register_buffer('attn_mask', text.attn_mask, persistent=False)
|
||||
|
||||
self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
|
||||
|
||||
def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
|
||||
# lock image tower as per LiT - https://arxiv.org/abs/2111.07991
|
||||
self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
|
||||
|
||||
@torch.jit.ignore
|
||||
def set_grad_checkpointing(self, enable=True):
|
||||
self.visual.set_grad_checkpointing(enable)
|
||||
self.transformer.grad_checkpointing = enable
|
||||
|
||||
@torch.jit.ignore
|
||||
def no_weight_decay(self):
|
||||
return {'logit_scale'}
|
||||
|
||||
def encode_image(self, image, normalize: bool = False):
|
||||
features = self.visual(image)
|
||||
return F.normalize(features, dim=-1) if normalize else features
|
||||
|
||||
def encode_text(self, text, normalize: bool = False):
|
||||
cast_dtype = self.transformer.get_cast_dtype()
|
||||
|
||||
x = self.token_embedding(text).to(cast_dtype) # [batch_size, n_ctx, d_model]
|
||||
|
||||
x = x + self.positional_embedding.to(cast_dtype)
|
||||
x = x.permute(1, 0, 2) # NLD -> LND
|
||||
x = self.transformer(x, attn_mask=self.attn_mask)
|
||||
x = x.permute(1, 0, 2) # LND -> NLD
|
||||
x = self.ln_final(x) # [batch_size, n_ctx, transformer.width]
|
||||
# take features from the eot embedding (eot_token is the highest number in each sequence)
|
||||
x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
|
||||
return F.normalize(x, dim=-1) if normalize else x
|
||||
|
||||
def forward(self, image, text):
|
||||
image_features = self.encode_image(image, normalize=True)
|
||||
text_features = self.encode_text(text, normalize=True)
|
||||
return image_features, text_features, self.logit_scale.exp()
|
||||
|
||||
|
||||
class CustomCLIP(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
embed_dim: int,
|
||||
vision_cfg: CLIPVisionCfg,
|
||||
text_cfg: CLIPTextCfg,
|
||||
quick_gelu: bool = False,
|
||||
cast_dtype: Optional[torch.dtype] = None,
|
||||
itm_task: bool = False,
|
||||
is_only_visual: bool = False,
|
||||
is_only_text: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
|
||||
self.text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
|
||||
self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) #可学习参数
|
||||
if is_only_visual:
|
||||
self.text = None
|
||||
if is_only_text:
|
||||
self.visual = None
|
||||
|
||||
def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
|
||||
# lock image tower as per LiT - https://arxiv.org/abs/2111.07991
|
||||
self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
|
||||
|
||||
def lock_text_tower(self, unlocked_layers:int=0, freeze_layer_norm:bool=True):
|
||||
self.text.lock(unlocked_layers, freeze_layer_norm)
|
||||
|
||||
@torch.jit.ignore
|
||||
def set_grad_checkpointing(self, enable=True):
|
||||
self.visual.set_grad_checkpointing(enable)
|
||||
if self.text is not None:
|
||||
self.text.set_grad_checkpointing(enable)
|
||||
|
||||
@torch.jit.ignore
|
||||
def no_weight_decay(self):
|
||||
return {'logit_scale'}
|
||||
|
||||
def encode_image(self, image, normalize: bool = False):
|
||||
features = self.visual(image)
|
||||
return F.normalize(features, dim=-1) if normalize else features
|
||||
|
||||
def encode_text(self, text, normalize: bool = False):
|
||||
features = self.text(text)
|
||||
return F.normalize(features, dim=-1) if normalize else features
|
||||
|
||||
def forward(self, image, text):
|
||||
if self.visual is not None:
|
||||
image_features = self.encode_image(image, normalize=True)
|
||||
else:
|
||||
image_features = None
|
||||
if self.text is not None:
|
||||
text_features = self.encode_text(text, normalize=True)
|
||||
else:
|
||||
text_features = None
|
||||
return image_features, text_features, self.logit_scale.exp()
|
||||
|
||||
|
||||
def convert_weights_to_lp(model: nn.Module, dtype=torch.float16):
|
||||
"""Convert applicable model parameters to low-precision (bf16 or fp16)"""
|
||||
|
||||
def _convert_weights(l):
|
||||
|
||||
if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
|
||||
l.weight.data = l.weight.data.to(dtype)
|
||||
if l.bias is not None:
|
||||
l.bias.data = l.bias.data.to(dtype)
|
||||
|
||||
if isinstance(l, (nn.MultiheadAttention, Attention)):
|
||||
for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
|
||||
tensor = getattr(l, attr, None)
|
||||
if tensor is not None:
|
||||
tensor.data = tensor.data.to(dtype)
|
||||
|
||||
if isinstance(l, nn.Parameter):
|
||||
l.data = l.data.to(dtype)
|
||||
|
||||
for name in ["text_projection", "proj"]:
|
||||
if hasattr(l, name) and isinstance(l, nn.Parameter):
|
||||
attr = getattr(l, name, None)
|
||||
if attr is not None:
|
||||
attr.data = attr.data.to(dtype)
|
||||
|
||||
model.apply(_convert_weights)
|
||||
|
||||
|
||||
convert_weights_to_fp16 = convert_weights_to_lp # backwards compat
|
||||
|
||||
|
||||
# used to maintain checkpoint compatibility
|
||||
def convert_to_custom_text_state_dict(state_dict: dict):
|
||||
if 'text_projection' in state_dict:
|
||||
# old format state_dict, move text tower -> .text
|
||||
new_state_dict = {}
|
||||
for k, v in state_dict.items():
|
||||
if any(k.startswith(p) for p in (
|
||||
'text_projection',
|
||||
'positional_embedding',
|
||||
'token_embedding',
|
||||
'transformer',
|
||||
'ln_final',
|
||||
'logit_scale'
|
||||
)):
|
||||
k = 'text.' + k
|
||||
new_state_dict[k] = v
|
||||
return new_state_dict
|
||||
return state_dict
|
||||
|
||||
|
||||
def build_model_from_openai_state_dict(
|
||||
state_dict: dict,
|
||||
quick_gelu=True,
|
||||
cast_dtype=torch.float16,
|
||||
):
|
||||
vit = "visual.proj" in state_dict
|
||||
|
||||
if vit:
|
||||
vision_width = state_dict["visual.conv1.weight"].shape[0]
|
||||
vision_layers = len(
|
||||
[k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
|
||||
vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
|
||||
grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
|
||||
image_size = vision_patch_size * grid_size
|
||||
else:
|
||||
counts: list = [
|
||||
len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
|
||||
vision_layers = tuple(counts)
|
||||
vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
|
||||
output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
|
||||
vision_patch_size = None
|
||||
assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
|
||||
image_size = output_width * 32
|
||||
|
||||
embed_dim = state_dict["text_projection"].shape[1]
|
||||
context_length = state_dict["positional_embedding"].shape[0]
|
||||
vocab_size = state_dict["token_embedding.weight"].shape[0]
|
||||
transformer_width = state_dict["ln_final.weight"].shape[0]
|
||||
transformer_heads = transformer_width // 64
|
||||
transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
|
||||
|
||||
vision_cfg = CLIPVisionCfg(
|
||||
layers=vision_layers,
|
||||
width=vision_width,
|
||||
patch_size=vision_patch_size,
|
||||
image_size=image_size,
|
||||
)
|
||||
text_cfg = CLIPTextCfg(
|
||||
context_length=context_length,
|
||||
vocab_size=vocab_size,
|
||||
width=transformer_width,
|
||||
heads=transformer_heads,
|
||||
layers=transformer_layers
|
||||
)
|
||||
model = CLIP(
|
||||
embed_dim,
|
||||
vision_cfg=vision_cfg,
|
||||
text_cfg=text_cfg,
|
||||
quick_gelu=quick_gelu, # OpenAI models were trained with QuickGELU
|
||||
cast_dtype=cast_dtype,
|
||||
)
|
||||
|
||||
for key in ["input_resolution", "context_length", "vocab_size"]:
|
||||
state_dict.pop(key, None)
|
||||
|
||||
convert_weights_to_fp16(model) # OpenAI state dicts are partially converted to float16
|
||||
model.load_state_dict(state_dict)
|
||||
return model.eval()
|
||||
|
||||
|
||||
def trace_model(model, batch_size=256, device=torch.device('cpu')):
|
||||
model.eval()
|
||||
image_size = model.visual.image_size
|
||||
example_images = torch.ones((batch_size, 3, image_size, image_size), device=device)
|
||||
example_text = torch.zeros((batch_size, model.context_length), dtype=torch.int, device=device)
|
||||
model = torch.jit.trace_module(
|
||||
model,
|
||||
inputs=dict(
|
||||
forward=(example_images, example_text),
|
||||
encode_text=(example_text,),
|
||||
encode_image=(example_images,)
|
||||
))
|
||||
model.visual.image_size = image_size
|
||||
return model
|
||||
@@ -0,0 +1,19 @@
|
||||
{
|
||||
"embed_dim": 512,
|
||||
"vision_cfg": {
|
||||
"image_size": 224,
|
||||
"layers": 12,
|
||||
"width": 768,
|
||||
"patch_size": 16,
|
||||
"eva_model_name": "eva-clip-b-16",
|
||||
"ls_init_value": 0.1,
|
||||
"drop_path_rate": 0.0
|
||||
},
|
||||
"text_cfg": {
|
||||
"context_length": 77,
|
||||
"vocab_size": 49408,
|
||||
"width": 512,
|
||||
"heads": 8,
|
||||
"layers": 12
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
{
|
||||
"embed_dim": 1024,
|
||||
"vision_cfg": {
|
||||
"image_size": 224,
|
||||
"layers": 40,
|
||||
"width": 1408,
|
||||
"head_width": 88,
|
||||
"mlp_ratio": 4.3637,
|
||||
"patch_size": 14,
|
||||
"eva_model_name": "eva-clip-g-14-x",
|
||||
"drop_path_rate": 0,
|
||||
"xattn": true,
|
||||
"fusedLN": true
|
||||
},
|
||||
"text_cfg": {
|
||||
"context_length": 77,
|
||||
"vocab_size": 49408,
|
||||
"width": 1024,
|
||||
"heads": 16,
|
||||
"layers": 24,
|
||||
"xattn": false,
|
||||
"fusedLN": true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
{
|
||||
"embed_dim": 1024,
|
||||
"vision_cfg": {
|
||||
"image_size": 224,
|
||||
"layers": 40,
|
||||
"width": 1408,
|
||||
"head_width": 88,
|
||||
"mlp_ratio": 4.3637,
|
||||
"patch_size": 14,
|
||||
"eva_model_name": "eva-clip-g-14-x",
|
||||
"drop_path_rate": 0.4,
|
||||
"xattn": true,
|
||||
"fusedLN": true
|
||||
},
|
||||
"text_cfg": {
|
||||
"context_length": 77,
|
||||
"vocab_size": 49408,
|
||||
"width": 768,
|
||||
"heads": 12,
|
||||
"layers": 12,
|
||||
"xattn": false,
|
||||
"fusedLN": true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
{
|
||||
"embed_dim": 512,
|
||||
"vision_cfg": {
|
||||
"image_size": 224,
|
||||
"layers": 12,
|
||||
"width": 768,
|
||||
"head_width": 64,
|
||||
"patch_size": 16,
|
||||
"mlp_ratio": 2.6667,
|
||||
"eva_model_name": "eva-clip-b-16-X",
|
||||
"drop_path_rate": 0.0,
|
||||
"xattn": true,
|
||||
"fusedLN": true,
|
||||
"rope": true,
|
||||
"pt_hw_seq_len": 16,
|
||||
"intp_freq": true,
|
||||
"naiveswiglu": true,
|
||||
"subln": true,
|
||||
"patch_dropout": 0.5
|
||||
},
|
||||
"text_cfg": {
|
||||
"context_length": 77,
|
||||
"vocab_size": 49408,
|
||||
"width": 512,
|
||||
"heads": 8,
|
||||
"layers": 12,
|
||||
"xattn": true,
|
||||
"fusedLN": true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
{
|
||||
"embed_dim": 768,
|
||||
"vision_cfg": {
|
||||
"image_size": 336,
|
||||
"layers": 24,
|
||||
"width": 1024,
|
||||
"drop_path_rate": 0,
|
||||
"head_width": 64,
|
||||
"mlp_ratio": 2.6667,
|
||||
"patch_size": 14,
|
||||
"eva_model_name": "eva-clip-l-14-336",
|
||||
"xattn": true,
|
||||
"fusedLN": true,
|
||||
"rope": true,
|
||||
"pt_hw_seq_len": 16,
|
||||
"intp_freq": true,
|
||||
"naiveswiglu": true,
|
||||
"subln": true
|
||||
},
|
||||
"text_cfg": {
|
||||
"context_length": 77,
|
||||
"vocab_size": 49408,
|
||||
"width": 768,
|
||||
"heads": 12,
|
||||
"layers": 12,
|
||||
"xattn": false,
|
||||
"fusedLN": true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
{
|
||||
"embed_dim": 768,
|
||||
"vision_cfg": {
|
||||
"image_size": 224,
|
||||
"layers": 24,
|
||||
"width": 1024,
|
||||
"drop_path_rate": 0,
|
||||
"head_width": 64,
|
||||
"mlp_ratio": 2.6667,
|
||||
"patch_size": 14,
|
||||
"eva_model_name": "eva-clip-l-14",
|
||||
"xattn": true,
|
||||
"fusedLN": true,
|
||||
"rope": true,
|
||||
"pt_hw_seq_len": 16,
|
||||
"intp_freq": true,
|
||||
"naiveswiglu": true,
|
||||
"subln": true
|
||||
},
|
||||
"text_cfg": {
|
||||
"context_length": 77,
|
||||
"vocab_size": 49408,
|
||||
"width": 768,
|
||||
"heads": 12,
|
||||
"layers": 12,
|
||||
"xattn": false,
|
||||
"fusedLN": true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
{
|
||||
"embed_dim": 1024,
|
||||
"vision_cfg": {
|
||||
"image_size": 224,
|
||||
"layers": 64,
|
||||
"width": 1792,
|
||||
"head_width": 112,
|
||||
"mlp_ratio": 8.571428571428571,
|
||||
"patch_size": 14,
|
||||
"eva_model_name": "eva-clip-4b-14-x",
|
||||
"drop_path_rate": 0,
|
||||
"xattn": true,
|
||||
"postnorm": true,
|
||||
"fusedLN": true
|
||||
},
|
||||
"text_cfg": {
|
||||
"context_length": 77,
|
||||
"vocab_size": 49408,
|
||||
"width": 1280,
|
||||
"heads": 20,
|
||||
"layers": 32,
|
||||
"xattn": false,
|
||||
"fusedLN": true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
{
|
||||
"embed_dim": 1024,
|
||||
"vision_cfg": {
|
||||
"image_size": 224,
|
||||
"layers": 64,
|
||||
"width": 1792,
|
||||
"head_width": 112,
|
||||
"mlp_ratio": 8.571428571428571,
|
||||
"patch_size": 14,
|
||||
"eva_model_name": "eva-clip-4b-14-x",
|
||||
"drop_path_rate": 0,
|
||||
"xattn": true,
|
||||
"postnorm": true,
|
||||
"fusedLN": true
|
||||
},
|
||||
"text_cfg": {
|
||||
"context_length": 77,
|
||||
"vocab_size": 49408,
|
||||
"width": 1024,
|
||||
"heads": 16,
|
||||
"layers": 24,
|
||||
"xattn": false,
|
||||
"fusedLN": true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
from collections import OrderedDict
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from .utils import freeze_batch_norm_2d
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1):
|
||||
super().__init__()
|
||||
|
||||
# all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
|
||||
self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(planes)
|
||||
self.act1 = nn.ReLU(inplace=True)
|
||||
|
||||
self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(planes)
|
||||
self.act2 = nn.ReLU(inplace=True)
|
||||
|
||||
self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
|
||||
|
||||
self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
|
||||
self.act3 = nn.ReLU(inplace=True)
|
||||
|
||||
self.downsample = None
|
||||
self.stride = stride
|
||||
|
||||
if stride > 1 or inplanes != planes * Bottleneck.expansion:
|
||||
# downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
|
||||
self.downsample = nn.Sequential(OrderedDict([
|
||||
("-1", nn.AvgPool2d(stride)),
|
||||
("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
|
||||
("1", nn.BatchNorm2d(planes * self.expansion))
|
||||
]))
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
identity = x
|
||||
|
||||
out = self.act1(self.bn1(self.conv1(x)))
|
||||
out = self.act2(self.bn2(self.conv2(out)))
|
||||
out = self.avgpool(out)
|
||||
out = self.bn3(self.conv3(out))
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.act3(out)
|
||||
return out
|
||||
|
||||
|
||||
class AttentionPool2d(nn.Module):
|
||||
def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
|
||||
super().__init__()
|
||||
self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
|
||||
self.k_proj = nn.Linear(embed_dim, embed_dim)
|
||||
self.q_proj = nn.Linear(embed_dim, embed_dim)
|
||||
self.v_proj = nn.Linear(embed_dim, embed_dim)
|
||||
self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
|
||||
self.num_heads = num_heads
|
||||
|
||||
def forward(self, x):
|
||||
x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1) # NCHW -> (HW)NC
|
||||
x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (HW+1)NC
|
||||
x = x + self.positional_embedding[:, None, :].to(x.dtype) # (HW+1)NC
|
||||
x, _ = F.multi_head_attention_forward(
|
||||
query=x, key=x, value=x,
|
||||
embed_dim_to_check=x.shape[-1],
|
||||
num_heads=self.num_heads,
|
||||
q_proj_weight=self.q_proj.weight,
|
||||
k_proj_weight=self.k_proj.weight,
|
||||
v_proj_weight=self.v_proj.weight,
|
||||
in_proj_weight=None,
|
||||
in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
|
||||
bias_k=None,
|
||||
bias_v=None,
|
||||
add_zero_attn=False,
|
||||
dropout_p=0.,
|
||||
out_proj_weight=self.c_proj.weight,
|
||||
out_proj_bias=self.c_proj.bias,
|
||||
use_separate_proj_weight=True,
|
||||
training=self.training,
|
||||
need_weights=False
|
||||
)
|
||||
|
||||
return x[0]
|
||||
|
||||
|
||||
class ModifiedResNet(nn.Module):
|
||||
"""
|
||||
A ResNet class that is similar to torchvision's but contains the following changes:
|
||||
- There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
|
||||
- Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
|
||||
- The final pooling layer is a QKV attention instead of an average pool
|
||||
"""
|
||||
|
||||
def __init__(self, layers, output_dim, heads, image_size=224, width=64):
|
||||
super().__init__()
|
||||
self.output_dim = output_dim
|
||||
self.image_size = image_size
|
||||
|
||||
# the 3-layer stem
|
||||
self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(width // 2)
|
||||
self.act1 = nn.ReLU(inplace=True)
|
||||
self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(width // 2)
|
||||
self.act2 = nn.ReLU(inplace=True)
|
||||
self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(width)
|
||||
self.act3 = nn.ReLU(inplace=True)
|
||||
self.avgpool = nn.AvgPool2d(2)
|
||||
|
||||
# residual layers
|
||||
self._inplanes = width # this is a *mutable* variable used during construction
|
||||
self.layer1 = self._make_layer(width, layers[0])
|
||||
self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
|
||||
self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
|
||||
|
||||
embed_dim = width * 32 # the ResNet feature dimension
|
||||
self.attnpool = AttentionPool2d(image_size // 32, embed_dim, heads, output_dim)
|
||||
|
||||
self.init_parameters()
|
||||
|
||||
def _make_layer(self, planes, blocks, stride=1):
|
||||
layers = [Bottleneck(self._inplanes, planes, stride)]
|
||||
|
||||
self._inplanes = planes * Bottleneck.expansion
|
||||
for _ in range(1, blocks):
|
||||
layers.append(Bottleneck(self._inplanes, planes))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def init_parameters(self):
|
||||
if self.attnpool is not None:
|
||||
std = self.attnpool.c_proj.in_features ** -0.5
|
||||
nn.init.normal_(self.attnpool.q_proj.weight, std=std)
|
||||
nn.init.normal_(self.attnpool.k_proj.weight, std=std)
|
||||
nn.init.normal_(self.attnpool.v_proj.weight, std=std)
|
||||
nn.init.normal_(self.attnpool.c_proj.weight, std=std)
|
||||
|
||||
for resnet_block in [self.layer1, self.layer2, self.layer3, self.layer4]:
|
||||
for name, param in resnet_block.named_parameters():
|
||||
if name.endswith("bn3.weight"):
|
||||
nn.init.zeros_(param)
|
||||
|
||||
def lock(self, unlocked_groups=0, freeze_bn_stats=False):
|
||||
assert unlocked_groups == 0, 'partial locking not currently supported for this model'
|
||||
for param in self.parameters():
|
||||
param.requires_grad = False
|
||||
if freeze_bn_stats:
|
||||
freeze_batch_norm_2d(self)
|
||||
|
||||
@torch.jit.ignore
|
||||
def set_grad_checkpointing(self, enable=True):
|
||||
# FIXME support for non-transformer
|
||||
pass
|
||||
|
||||
def stem(self, x):
|
||||
x = self.act1(self.bn1(self.conv1(x)))
|
||||
x = self.act2(self.bn2(self.conv2(x)))
|
||||
x = self.act3(self.bn3(self.conv3(x)))
|
||||
x = self.avgpool(x)
|
||||
return x
|
||||
|
||||
def forward(self, x):
|
||||
x = self.stem(x)
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
x = self.layer4(x)
|
||||
x = self.attnpool(x)
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,144 @@
|
||||
""" OpenAI pretrained model functions
|
||||
|
||||
Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
|
||||
"""
|
||||
|
||||
import os
|
||||
import warnings
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import torch
|
||||
|
||||
from .model import build_model_from_openai_state_dict, convert_weights_to_lp, get_cast_dtype
|
||||
from .pretrained import get_pretrained_url, list_pretrained_models_by_tag, download_pretrained_from_url
|
||||
|
||||
__all__ = ["list_openai_models", "load_openai_model"]
|
||||
|
||||
|
||||
def list_openai_models() -> List[str]:
|
||||
"""Returns the names of available CLIP models"""
|
||||
return list_pretrained_models_by_tag('openai')
|
||||
|
||||
|
||||
def load_openai_model(
|
||||
name: str,
|
||||
precision: Optional[str] = None,
|
||||
device: Optional[Union[str, torch.device]] = None,
|
||||
jit: bool = True,
|
||||
cache_dir: Optional[str] = None,
|
||||
):
|
||||
"""Load a CLIP model
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
|
||||
precision: str
|
||||
Model precision, if None defaults to 'fp32' if device == 'cpu' else 'fp16'.
|
||||
device : Union[str, torch.device]
|
||||
The device to put the loaded model
|
||||
jit : bool
|
||||
Whether to load the optimized JIT model (default) or more hackable non-JIT model.
|
||||
cache_dir : Optional[str]
|
||||
The directory to cache the downloaded model weights
|
||||
|
||||
Returns
|
||||
-------
|
||||
model : torch.nn.Module
|
||||
The CLIP model
|
||||
preprocess : Callable[[PIL.Image], torch.Tensor]
|
||||
A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
|
||||
"""
|
||||
if device is None:
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
if precision is None:
|
||||
precision = 'fp32' if device == 'cpu' else 'fp16'
|
||||
|
||||
if get_pretrained_url(name, 'openai'):
|
||||
model_path = download_pretrained_from_url(get_pretrained_url(name, 'openai'), cache_dir=cache_dir)
|
||||
elif os.path.isfile(name):
|
||||
model_path = name
|
||||
else:
|
||||
raise RuntimeError(f"Model {name} not found; available models = {list_openai_models()}")
|
||||
|
||||
try:
|
||||
# loading JIT archive
|
||||
model = torch.jit.load(model_path, map_location=device if jit else "cpu").eval()
|
||||
state_dict = None
|
||||
except RuntimeError:
|
||||
# loading saved state dict
|
||||
if jit:
|
||||
warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
|
||||
jit = False
|
||||
state_dict = torch.load(model_path, map_location="cpu")
|
||||
|
||||
if not jit:
|
||||
# Build a non-jit model from the OpenAI jitted model state dict
|
||||
cast_dtype = get_cast_dtype(precision)
|
||||
try:
|
||||
model = build_model_from_openai_state_dict(state_dict or model.state_dict(), cast_dtype=cast_dtype)
|
||||
except KeyError:
|
||||
sd = {k[7:]: v for k, v in state_dict["state_dict"].items()}
|
||||
model = build_model_from_openai_state_dict(sd, cast_dtype=cast_dtype)
|
||||
|
||||
# model from OpenAI state dict is in manually cast fp16 mode, must be converted for AMP/fp32/bf16 use
|
||||
model = model.to(device)
|
||||
if precision.startswith('amp') or precision == 'fp32':
|
||||
model.float()
|
||||
elif precision == 'bf16':
|
||||
convert_weights_to_lp(model, dtype=torch.bfloat16)
|
||||
|
||||
return model
|
||||
|
||||
# patch the device names
|
||||
device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
|
||||
device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
|
||||
|
||||
def patch_device(module):
|
||||
try:
|
||||
graphs = [module.graph] if hasattr(module, "graph") else []
|
||||
except RuntimeError:
|
||||
graphs = []
|
||||
|
||||
if hasattr(module, "forward1"):
|
||||
graphs.append(module.forward1.graph)
|
||||
|
||||
for graph in graphs:
|
||||
for node in graph.findAllNodes("prim::Constant"):
|
||||
if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"):
|
||||
node.copyAttributes(device_node)
|
||||
|
||||
model.apply(patch_device)
|
||||
patch_device(model.encode_image)
|
||||
patch_device(model.encode_text)
|
||||
|
||||
# patch dtype to float32 (typically for CPU)
|
||||
if precision == 'fp32':
|
||||
float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
|
||||
float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
|
||||
float_node = float_input.node()
|
||||
|
||||
def patch_float(module):
|
||||
try:
|
||||
graphs = [module.graph] if hasattr(module, "graph") else []
|
||||
except RuntimeError:
|
||||
graphs = []
|
||||
|
||||
if hasattr(module, "forward1"):
|
||||
graphs.append(module.forward1.graph)
|
||||
|
||||
for graph in graphs:
|
||||
for node in graph.findAllNodes("aten::to"):
|
||||
inputs = list(node.inputs())
|
||||
for i in [1, 2]: # dtype can be the second or third argument to aten::to()
|
||||
if inputs[i].node()["value"] == 5:
|
||||
inputs[i].node().copyAttributes(float_node)
|
||||
|
||||
model.apply(patch_float)
|
||||
patch_float(model.encode_image)
|
||||
patch_float(model.encode_text)
|
||||
model.float()
|
||||
|
||||
# ensure image_size attr available at consistent location for both jit and non-jit
|
||||
model.visual.image_size = model.input_resolution.item()
|
||||
return model
|
||||
@@ -0,0 +1,332 @@
|
||||
import hashlib
|
||||
import os
|
||||
import urllib
|
||||
import warnings
|
||||
from functools import partial
|
||||
from typing import Dict, Union
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
_has_hf_hub = True
|
||||
except ImportError:
|
||||
hf_hub_download = None
|
||||
_has_hf_hub = False
|
||||
|
||||
|
||||
def _pcfg(url='', hf_hub='', filename='', mean=None, std=None):
|
||||
return dict(
|
||||
url=url,
|
||||
hf_hub=hf_hub,
|
||||
mean=mean,
|
||||
std=std,
|
||||
)
|
||||
|
||||
_VITB32 = dict(
|
||||
openai=_pcfg(
|
||||
"https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt"),
|
||||
laion400m_e31=_pcfg(
|
||||
"https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e31-d867053b.pt"),
|
||||
laion400m_e32=_pcfg(
|
||||
"https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e32-46683a32.pt"),
|
||||
laion2b_e16=_pcfg(
|
||||
"https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-laion2b_e16-af8dbd0c.pth"),
|
||||
laion2b_s34b_b79k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-laion2B-s34B-b79K/')
|
||||
)
|
||||
|
||||
_VITB32_quickgelu = dict(
|
||||
openai=_pcfg(
|
||||
"https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt"),
|
||||
laion400m_e31=_pcfg(
|
||||
"https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e31-d867053b.pt"),
|
||||
laion400m_e32=_pcfg(
|
||||
"https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e32-46683a32.pt"),
|
||||
)
|
||||
|
||||
_VITB16 = dict(
|
||||
openai=_pcfg(
|
||||
"https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt"),
|
||||
laion400m_e31=_pcfg(
|
||||
"https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16-laion400m_e31-00efa78f.pt"),
|
||||
laion400m_e32=_pcfg(
|
||||
"https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16-laion400m_e32-55e67d44.pt"),
|
||||
laion2b_s34b_b88k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-laion2B-s34B-b88K/'),
|
||||
)
|
||||
|
||||
_EVAB16 = dict(
|
||||
eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_B_psz14to16.pt'),
|
||||
eva02=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_B_psz14to16.pt'),
|
||||
eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_B_psz16_s8B.pt'),
|
||||
eva02_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_B_psz16_s8B.pt'),
|
||||
)
|
||||
|
||||
_VITB16_PLUS_240 = dict(
|
||||
laion400m_e31=_pcfg(
|
||||
"https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16_plus_240-laion400m_e31-8fb26589.pt"),
|
||||
laion400m_e32=_pcfg(
|
||||
"https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16_plus_240-laion400m_e32-699c4b84.pt"),
|
||||
)
|
||||
|
||||
_VITL14 = dict(
|
||||
openai=_pcfg(
|
||||
"https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt"),
|
||||
laion400m_e31=_pcfg(
|
||||
"https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_l_14-laion400m_e31-69988bb6.pt"),
|
||||
laion400m_e32=_pcfg(
|
||||
"https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_l_14-laion400m_e32-3d133497.pt"),
|
||||
laion2b_s32b_b82k=_pcfg(
|
||||
hf_hub='laion/CLIP-ViT-L-14-laion2B-s32B-b82K/',
|
||||
mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
|
||||
)
|
||||
|
||||
_EVAL14 = dict(
|
||||
eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_L_psz14.pt'),
|
||||
eva02=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_L_psz14.pt'),
|
||||
eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_psz14_s4B.pt'),
|
||||
eva02_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_psz14_s4B.pt'),
|
||||
)
|
||||
|
||||
_VITL14_336 = dict(
|
||||
openai=_pcfg(
|
||||
"https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt"),
|
||||
)
|
||||
|
||||
_EVAL14_336 = dict(
|
||||
eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_336_psz14_s6B.pt'),
|
||||
eva02_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_336_psz14_s6B.pt'),
|
||||
eva_clip_224to336=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_psz14_224to336.pt'),
|
||||
eva02_clip_224to336=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_psz14_224to336.pt'),
|
||||
)
|
||||
|
||||
_VITH14 = dict(
|
||||
laion2b_s32b_b79k=_pcfg(hf_hub='laion/CLIP-ViT-H-14-laion2B-s32B-b79K/'),
|
||||
)
|
||||
|
||||
_VITg14 = dict(
|
||||
laion2b_s12b_b42k=_pcfg(hf_hub='laion/CLIP-ViT-g-14-laion2B-s12B-b42K/'),
|
||||
laion2b_s34b_b88k=_pcfg(hf_hub='laion/CLIP-ViT-g-14-laion2B-s34B-b88K/'),
|
||||
)
|
||||
|
||||
_EVAg14 = dict(
|
||||
eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/'),
|
||||
eva01=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_g_psz14.pt'),
|
||||
eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_CLIP_g_14_psz14_s11B.pt'),
|
||||
eva01_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_CLIP_g_14_psz14_s11B.pt'),
|
||||
)
|
||||
|
||||
_EVAg14_PLUS = dict(
|
||||
eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/'),
|
||||
eva01=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_g_psz14.pt'),
|
||||
eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_CLIP_g_14_plus_psz14_s11B.pt'),
|
||||
eva01_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_CLIP_g_14_plus_psz14_s11B.pt'),
|
||||
)
|
||||
|
||||
_VITbigG14 = dict(
|
||||
laion2b_s39b_b160k=_pcfg(hf_hub='laion/CLIP-ViT-bigG-14-laion2B-39B-b160k/'),
|
||||
)
|
||||
|
||||
_EVAbigE14 = dict(
|
||||
eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_E_psz14.pt'),
|
||||
eva02=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_E_psz14.pt'),
|
||||
eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_E_psz14_s4B.pt'),
|
||||
eva02_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_E_psz14_s4B.pt'),
|
||||
)
|
||||
|
||||
_EVAbigE14_PLUS = dict(
|
||||
eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_E_psz14.pt'),
|
||||
eva02=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_E_psz14.pt'),
|
||||
eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_E_psz14_plus_s9B.pt'),
|
||||
eva02_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_E_psz14_plus_s9B.pt'),
|
||||
)
|
||||
|
||||
|
||||
_PRETRAINED = {
|
||||
# "ViT-B-32": _VITB32,
|
||||
"OpenaiCLIP-B-32": _VITB32,
|
||||
"OpenCLIP-B-32": _VITB32,
|
||||
|
||||
# "ViT-B-32-quickgelu": _VITB32_quickgelu,
|
||||
"OpenaiCLIP-B-32-quickgelu": _VITB32_quickgelu,
|
||||
"OpenCLIP-B-32-quickgelu": _VITB32_quickgelu,
|
||||
|
||||
# "ViT-B-16": _VITB16,
|
||||
"OpenaiCLIP-B-16": _VITB16,
|
||||
"OpenCLIP-B-16": _VITB16,
|
||||
|
||||
"EVA02-B-16": _EVAB16,
|
||||
"EVA02-CLIP-B-16": _EVAB16,
|
||||
|
||||
# "ViT-B-16-plus-240": _VITB16_PLUS_240,
|
||||
"OpenCLIP-B-16-plus-240": _VITB16_PLUS_240,
|
||||
|
||||
# "ViT-L-14": _VITL14,
|
||||
"OpenaiCLIP-L-14": _VITL14,
|
||||
"OpenCLIP-L-14": _VITL14,
|
||||
|
||||
"EVA02-L-14": _EVAL14,
|
||||
"EVA02-CLIP-L-14": _EVAL14,
|
||||
|
||||
# "ViT-L-14-336": _VITL14_336,
|
||||
"OpenaiCLIP-L-14-336": _VITL14_336,
|
||||
|
||||
"EVA02-CLIP-L-14-336": _EVAL14_336,
|
||||
|
||||
# "ViT-H-14": _VITH14,
|
||||
# "ViT-g-14": _VITg14,
|
||||
"OpenCLIP-H-14": _VITH14,
|
||||
"OpenCLIP-g-14": _VITg14,
|
||||
|
||||
"EVA01-CLIP-g-14": _EVAg14,
|
||||
"EVA01-CLIP-g-14-plus": _EVAg14_PLUS,
|
||||
|
||||
# "ViT-bigG-14": _VITbigG14,
|
||||
"OpenCLIP-bigG-14": _VITbigG14,
|
||||
|
||||
"EVA02-CLIP-bigE-14": _EVAbigE14,
|
||||
"EVA02-CLIP-bigE-14-plus": _EVAbigE14_PLUS,
|
||||
}
|
||||
|
||||
|
||||
def _clean_tag(tag: str):
|
||||
# normalize pretrained tags
|
||||
return tag.lower().replace('-', '_')
|
||||
|
||||
|
||||
def list_pretrained(as_str: bool = False):
|
||||
""" returns list of pretrained models
|
||||
Returns a tuple (model_name, pretrain_tag) by default or 'name:tag' if as_str == True
|
||||
"""
|
||||
return [':'.join([k, t]) if as_str else (k, t) for k in _PRETRAINED.keys() for t in _PRETRAINED[k].keys()]
|
||||
|
||||
|
||||
def list_pretrained_models_by_tag(tag: str):
|
||||
""" return all models having the specified pretrain tag """
|
||||
models = []
|
||||
tag = _clean_tag(tag)
|
||||
for k in _PRETRAINED.keys():
|
||||
if tag in _PRETRAINED[k]:
|
||||
models.append(k)
|
||||
return models
|
||||
|
||||
|
||||
def list_pretrained_tags_by_model(model: str):
|
||||
""" return all pretrain tags for the specified model architecture """
|
||||
tags = []
|
||||
if model in _PRETRAINED:
|
||||
tags.extend(_PRETRAINED[model].keys())
|
||||
return tags
|
||||
|
||||
|
||||
def is_pretrained_cfg(model: str, tag: str):
|
||||
if model not in _PRETRAINED:
|
||||
return False
|
||||
return _clean_tag(tag) in _PRETRAINED[model]
|
||||
|
||||
|
||||
def get_pretrained_cfg(model: str, tag: str):
|
||||
if model not in _PRETRAINED:
|
||||
return {}
|
||||
model_pretrained = _PRETRAINED[model]
|
||||
return model_pretrained.get(_clean_tag(tag), {})
|
||||
|
||||
|
||||
def get_pretrained_url(model: str, tag: str):
|
||||
cfg = get_pretrained_cfg(model, _clean_tag(tag))
|
||||
return cfg.get('url', '')
|
||||
|
||||
|
||||
def download_pretrained_from_url(
|
||||
url: str,
|
||||
cache_dir: Union[str, None] = None,
|
||||
):
|
||||
if not cache_dir:
|
||||
cache_dir = os.path.expanduser("~/.cache/clip")
|
||||
os.makedirs(cache_dir, exist_ok=True)
|
||||
filename = os.path.basename(url)
|
||||
|
||||
if 'openaipublic' in url:
|
||||
expected_sha256 = url.split("/")[-2]
|
||||
elif 'mlfoundations' in url:
|
||||
expected_sha256 = os.path.splitext(filename)[0].split("-")[-1]
|
||||
else:
|
||||
expected_sha256 = ''
|
||||
|
||||
download_target = os.path.join(cache_dir, filename)
|
||||
|
||||
if os.path.exists(download_target) and not os.path.isfile(download_target):
|
||||
raise RuntimeError(f"{download_target} exists and is not a regular file")
|
||||
|
||||
if os.path.isfile(download_target):
|
||||
if expected_sha256:
|
||||
if hashlib.sha256(open(download_target, "rb").read()).hexdigest().startswith(expected_sha256):
|
||||
return download_target
|
||||
else:
|
||||
warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
|
||||
else:
|
||||
return download_target
|
||||
|
||||
with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
|
||||
with tqdm(total=int(source.headers.get("Content-Length")), ncols=80, unit='iB', unit_scale=True) as loop:
|
||||
while True:
|
||||
buffer = source.read(8192)
|
||||
if not buffer:
|
||||
break
|
||||
|
||||
output.write(buffer)
|
||||
loop.update(len(buffer))
|
||||
|
||||
if expected_sha256 and not hashlib.sha256(open(download_target, "rb").read()).hexdigest().startswith(expected_sha256):
|
||||
raise RuntimeError(f"Model has been downloaded but the SHA256 checksum does not not match")
|
||||
|
||||
return download_target
|
||||
|
||||
|
||||
def has_hf_hub(necessary=False):
|
||||
if not _has_hf_hub and necessary:
|
||||
# if no HF Hub module installed, and it is necessary to continue, raise error
|
||||
raise RuntimeError(
|
||||
'Hugging Face hub model specified but package not installed. Run `pip install huggingface_hub`.')
|
||||
return _has_hf_hub
|
||||
|
||||
|
||||
def download_pretrained_from_hf(
|
||||
model_id: str,
|
||||
filename: str = 'open_clip_pytorch_model.bin',
|
||||
revision=None,
|
||||
cache_dir: Union[str, None] = None,
|
||||
):
|
||||
has_hf_hub(True)
|
||||
cached_file = hf_hub_download(model_id, filename, revision=revision, cache_dir=cache_dir)
|
||||
return cached_file
|
||||
|
||||
|
||||
def download_pretrained(
|
||||
cfg: Dict,
|
||||
force_hf_hub: bool = False,
|
||||
cache_dir: Union[str, None] = None,
|
||||
):
|
||||
target = ''
|
||||
if not cfg:
|
||||
return target
|
||||
|
||||
download_url = cfg.get('url', '')
|
||||
download_hf_hub = cfg.get('hf_hub', '')
|
||||
if download_hf_hub and force_hf_hub:
|
||||
# use HF hub even if url exists
|
||||
download_url = ''
|
||||
|
||||
if download_url:
|
||||
target = download_pretrained_from_url(download_url, cache_dir=cache_dir)
|
||||
elif download_hf_hub:
|
||||
has_hf_hub(True)
|
||||
# we assume the hf_hub entries in pretrained config combine model_id + filename in
|
||||
# 'org/model_name/filename.pt' form. To specify just the model id w/o filename and
|
||||
# use 'open_clip_pytorch_model.bin' default, there must be a trailing slash 'org/model_name/'.
|
||||
model_id, filename = os.path.split(download_hf_hub)
|
||||
if filename:
|
||||
target = download_pretrained_from_hf(model_id, filename=filename, cache_dir=cache_dir)
|
||||
else:
|
||||
target = download_pretrained_from_hf(model_id, cache_dir=cache_dir)
|
||||
|
||||
return target
|
||||
@@ -0,0 +1,137 @@
|
||||
from math import pi
|
||||
import torch
|
||||
from torch import nn
|
||||
from einops import rearrange, repeat
|
||||
import logging
|
||||
|
||||
def broadcat(tensors, dim = -1):
|
||||
num_tensors = len(tensors)
|
||||
shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
|
||||
assert len(shape_lens) == 1, 'tensors must all have the same number of dimensions'
|
||||
shape_len = list(shape_lens)[0]
|
||||
dim = (dim + shape_len) if dim < 0 else dim
|
||||
dims = list(zip(*map(lambda t: list(t.shape), tensors)))
|
||||
expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
|
||||
assert all([*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]), 'invalid dimensions for broadcastable concatentation'
|
||||
max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))
|
||||
expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))
|
||||
expanded_dims.insert(dim, (dim, dims[dim]))
|
||||
expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))
|
||||
tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))
|
||||
return torch.cat(tensors, dim = dim)
|
||||
|
||||
def rotate_half(x):
|
||||
x = rearrange(x, '... (d r) -> ... d r', r = 2)
|
||||
x1, x2 = x.unbind(dim = -1)
|
||||
x = torch.stack((-x2, x1), dim = -1)
|
||||
return rearrange(x, '... d r -> ... (d r)')
|
||||
|
||||
|
||||
class VisionRotaryEmbedding(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
pt_seq_len,
|
||||
ft_seq_len=None,
|
||||
custom_freqs = None,
|
||||
freqs_for = 'lang',
|
||||
theta = 10000,
|
||||
max_freq = 10,
|
||||
num_freqs = 1,
|
||||
):
|
||||
super().__init__()
|
||||
if custom_freqs:
|
||||
freqs = custom_freqs
|
||||
elif freqs_for == 'lang':
|
||||
freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
|
||||
elif freqs_for == 'pixel':
|
||||
freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
|
||||
elif freqs_for == 'constant':
|
||||
freqs = torch.ones(num_freqs).float()
|
||||
else:
|
||||
raise ValueError(f'unknown modality {freqs_for}')
|
||||
|
||||
if ft_seq_len is None: ft_seq_len = pt_seq_len
|
||||
t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
|
||||
|
||||
freqs_h = torch.einsum('..., f -> ... f', t, freqs)
|
||||
freqs_h = repeat(freqs_h, '... n -> ... (n r)', r = 2)
|
||||
|
||||
freqs_w = torch.einsum('..., f -> ... f', t, freqs)
|
||||
freqs_w = repeat(freqs_w, '... n -> ... (n r)', r = 2)
|
||||
|
||||
freqs = broadcat((freqs_h[:, None, :], freqs_w[None, :, :]), dim = -1)
|
||||
|
||||
self.register_buffer("freqs_cos", freqs.cos())
|
||||
self.register_buffer("freqs_sin", freqs.sin())
|
||||
|
||||
logging.info(f'Shape of rope freq: {self.freqs_cos.shape}')
|
||||
|
||||
def forward(self, t, start_index = 0):
|
||||
rot_dim = self.freqs_cos.shape[-1]
|
||||
end_index = start_index + rot_dim
|
||||
assert rot_dim <= t.shape[-1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}'
|
||||
t_left, t, t_right = t[..., :start_index], t[..., start_index:end_index], t[..., end_index:]
|
||||
t = (t * self.freqs_cos) + (rotate_half(t) * self.freqs_sin)
|
||||
|
||||
return torch.cat((t_left, t, t_right), dim = -1)
|
||||
|
||||
class VisionRotaryEmbeddingFast(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
pt_seq_len,
|
||||
ft_seq_len=None,
|
||||
custom_freqs = None,
|
||||
freqs_for = 'lang',
|
||||
theta = 10000,
|
||||
max_freq = 10,
|
||||
num_freqs = 1,
|
||||
patch_dropout = 0.
|
||||
):
|
||||
super().__init__()
|
||||
if custom_freqs:
|
||||
freqs = custom_freqs
|
||||
elif freqs_for == 'lang':
|
||||
freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
|
||||
elif freqs_for == 'pixel':
|
||||
freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
|
||||
elif freqs_for == 'constant':
|
||||
freqs = torch.ones(num_freqs).float()
|
||||
else:
|
||||
raise ValueError(f'unknown modality {freqs_for}')
|
||||
|
||||
if ft_seq_len is None: ft_seq_len = pt_seq_len
|
||||
t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
|
||||
|
||||
freqs = torch.einsum('..., f -> ... f', t, freqs)
|
||||
freqs = repeat(freqs, '... n -> ... (n r)', r = 2)
|
||||
freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim = -1)
|
||||
|
||||
freqs_cos = freqs.cos().view(-1, freqs.shape[-1])
|
||||
freqs_sin = freqs.sin().view(-1, freqs.shape[-1])
|
||||
|
||||
self.patch_dropout = patch_dropout
|
||||
|
||||
self.register_buffer("freqs_cos", freqs_cos)
|
||||
self.register_buffer("freqs_sin", freqs_sin)
|
||||
|
||||
logging.info(f'Shape of rope freq: {self.freqs_cos.shape}')
|
||||
|
||||
def forward(self, t, patch_indices_keep=None):
|
||||
if patch_indices_keep is not None:
|
||||
batch = t.size()[0]
|
||||
batch_indices = torch.arange(batch)
|
||||
batch_indices = batch_indices[..., None]
|
||||
|
||||
freqs_cos = repeat(self.freqs_cos, 'i j -> n i m j', n=t.shape[0], m=t.shape[1])
|
||||
freqs_sin = repeat(self.freqs_sin, 'i j -> n i m j', n=t.shape[0], m=t.shape[1])
|
||||
|
||||
freqs_cos = freqs_cos[batch_indices, patch_indices_keep]
|
||||
freqs_cos = rearrange(freqs_cos, 'n i m j -> n m i j')
|
||||
freqs_sin = freqs_sin[batch_indices, patch_indices_keep]
|
||||
freqs_sin = rearrange(freqs_sin, 'n i m j -> n m i j')
|
||||
|
||||
return t * freqs_cos + rotate_half(t) * freqs_sin
|
||||
|
||||
return t * self.freqs_cos + rotate_half(t) * self.freqs_sin
|
||||
@@ -0,0 +1,122 @@
|
||||
""" timm model adapter
|
||||
|
||||
Wraps timm (https://github.com/rwightman/pytorch-image-models) models for use as a vision tower in CLIP model.
|
||||
"""
|
||||
import logging
|
||||
from collections import OrderedDict
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
try:
|
||||
import timm
|
||||
from timm.models.layers import Mlp, to_2tuple
|
||||
try:
|
||||
# old timm imports < 0.8.1
|
||||
from timm.models.layers.attention_pool2d import RotAttentionPool2d
|
||||
from timm.models.layers.attention_pool2d import AttentionPool2d as AbsAttentionPool2d
|
||||
except ImportError:
|
||||
# new timm imports >= 0.8.1
|
||||
from timm.layers import RotAttentionPool2d
|
||||
from timm.layers import AttentionPool2d as AbsAttentionPool2d
|
||||
except ImportError:
|
||||
timm = None
|
||||
|
||||
from .utils import freeze_batch_norm_2d
|
||||
|
||||
|
||||
class TimmModel(nn.Module):
|
||||
""" timm model adapter
|
||||
# FIXME this adapter is a work in progress, may change in ways that break weight compat
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name,
|
||||
embed_dim,
|
||||
image_size=224,
|
||||
pool='avg',
|
||||
proj='linear',
|
||||
proj_bias=False,
|
||||
drop=0.,
|
||||
pretrained=False):
|
||||
super().__init__()
|
||||
if timm is None:
|
||||
raise RuntimeError("Please `pip install timm` to use timm models.")
|
||||
|
||||
self.image_size = to_2tuple(image_size)
|
||||
self.trunk = timm.create_model(model_name, pretrained=pretrained)
|
||||
feat_size = self.trunk.default_cfg.get('pool_size', None)
|
||||
feature_ndim = 1 if not feat_size else 2
|
||||
if pool in ('abs_attn', 'rot_attn'):
|
||||
assert feature_ndim == 2
|
||||
# if attn pooling used, remove both classifier and default pool
|
||||
self.trunk.reset_classifier(0, global_pool='')
|
||||
else:
|
||||
# reset global pool if pool config set, otherwise leave as network default
|
||||
reset_kwargs = dict(global_pool=pool) if pool else {}
|
||||
self.trunk.reset_classifier(0, **reset_kwargs)
|
||||
prev_chs = self.trunk.num_features
|
||||
|
||||
head_layers = OrderedDict()
|
||||
if pool == 'abs_attn':
|
||||
head_layers['pool'] = AbsAttentionPool2d(prev_chs, feat_size=feat_size, out_features=embed_dim)
|
||||
prev_chs = embed_dim
|
||||
elif pool == 'rot_attn':
|
||||
head_layers['pool'] = RotAttentionPool2d(prev_chs, out_features=embed_dim)
|
||||
prev_chs = embed_dim
|
||||
else:
|
||||
assert proj, 'projection layer needed if non-attention pooling is used.'
|
||||
|
||||
# NOTE attention pool ends with a projection layer, so proj should usually be set to '' if such pooling is used
|
||||
if proj == 'linear':
|
||||
head_layers['drop'] = nn.Dropout(drop)
|
||||
head_layers['proj'] = nn.Linear(prev_chs, embed_dim, bias=proj_bias)
|
||||
elif proj == 'mlp':
|
||||
head_layers['mlp'] = Mlp(prev_chs, 2 * embed_dim, embed_dim, drop=drop, bias=(True, proj_bias))
|
||||
|
||||
self.head = nn.Sequential(head_layers)
|
||||
|
||||
def lock(self, unlocked_groups=0, freeze_bn_stats=False):
|
||||
""" lock modules
|
||||
Args:
|
||||
unlocked_groups (int): leave last n layer groups unlocked (default: 0)
|
||||
"""
|
||||
if not unlocked_groups:
|
||||
# lock full model
|
||||
for param in self.trunk.parameters():
|
||||
param.requires_grad = False
|
||||
if freeze_bn_stats:
|
||||
freeze_batch_norm_2d(self.trunk)
|
||||
else:
|
||||
# NOTE: partial freeze requires latest timm (master) branch and is subject to change
|
||||
try:
|
||||
# FIXME import here until API stable and in an official release
|
||||
from timm.models.helpers import group_parameters, group_modules
|
||||
except ImportError:
|
||||
raise RuntimeError(
|
||||
'Please install latest timm `pip install git+https://github.com/rwightman/pytorch-image-models`')
|
||||
matcher = self.trunk.group_matcher()
|
||||
gparams = group_parameters(self.trunk, matcher)
|
||||
max_layer_id = max(gparams.keys())
|
||||
max_layer_id = max_layer_id - unlocked_groups
|
||||
for group_idx in range(max_layer_id + 1):
|
||||
group = gparams[group_idx]
|
||||
for param in group:
|
||||
self.trunk.get_parameter(param).requires_grad = False
|
||||
if freeze_bn_stats:
|
||||
gmodules = group_modules(self.trunk, matcher, reverse=True)
|
||||
gmodules = {k for k, v in gmodules.items() if v <= max_layer_id}
|
||||
freeze_batch_norm_2d(self.trunk, gmodules)
|
||||
|
||||
@torch.jit.ignore
|
||||
def set_grad_checkpointing(self, enable=True):
|
||||
try:
|
||||
self.trunk.set_grad_checkpointing(enable)
|
||||
except Exception as e:
|
||||
logging.warning('grad checkpointing not supported for this timm image tower, continuing without...')
|
||||
|
||||
def forward(self, x):
|
||||
x = self.trunk(x)
|
||||
x = self.head(x)
|
||||
return x
|
||||
@@ -0,0 +1,201 @@
|
||||
""" CLIP tokenizer
|
||||
|
||||
Copied from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
|
||||
"""
|
||||
import gzip
|
||||
import html
|
||||
import os
|
||||
from functools import lru_cache
|
||||
from typing import Union, List
|
||||
|
||||
import ftfy
|
||||
import regex as re
|
||||
import torch
|
||||
|
||||
# https://stackoverflow.com/q/62691279
|
||||
import os
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def default_bpe():
|
||||
return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def bytes_to_unicode():
|
||||
"""
|
||||
Returns list of utf-8 byte and a corresponding list of unicode strings.
|
||||
The reversible bpe codes work on unicode strings.
|
||||
This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
|
||||
When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
|
||||
This is a signficant percentage of your normal, say, 32K bpe vocab.
|
||||
To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
|
||||
And avoids mapping to whitespace/control characters the bpe code barfs on.
|
||||
"""
|
||||
bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
|
||||
cs = bs[:]
|
||||
n = 0
|
||||
for b in range(2**8):
|
||||
if b not in bs:
|
||||
bs.append(b)
|
||||
cs.append(2**8+n)
|
||||
n += 1
|
||||
cs = [chr(n) for n in cs]
|
||||
return dict(zip(bs, cs))
|
||||
|
||||
|
||||
def get_pairs(word):
|
||||
"""Return set of symbol pairs in a word.
|
||||
Word is represented as tuple of symbols (symbols being variable-length strings).
|
||||
"""
|
||||
pairs = set()
|
||||
prev_char = word[0]
|
||||
for char in word[1:]:
|
||||
pairs.add((prev_char, char))
|
||||
prev_char = char
|
||||
return pairs
|
||||
|
||||
|
||||
def basic_clean(text):
|
||||
text = ftfy.fix_text(text)
|
||||
text = html.unescape(html.unescape(text))
|
||||
return text.strip()
|
||||
|
||||
|
||||
def whitespace_clean(text):
|
||||
text = re.sub(r'\s+', ' ', text)
|
||||
text = text.strip()
|
||||
return text
|
||||
|
||||
|
||||
class SimpleTokenizer(object):
|
||||
def __init__(self, bpe_path: str = default_bpe(), special_tokens=None):
|
||||
self.byte_encoder = bytes_to_unicode()
|
||||
self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
|
||||
merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
|
||||
merges = merges[1:49152-256-2+1]
|
||||
merges = [tuple(merge.split()) for merge in merges]
|
||||
vocab = list(bytes_to_unicode().values())
|
||||
vocab = vocab + [v+'</w>' for v in vocab]
|
||||
for merge in merges:
|
||||
vocab.append(''.join(merge))
|
||||
if not special_tokens:
|
||||
special_tokens = ['<start_of_text>', '<end_of_text>']
|
||||
else:
|
||||
special_tokens = ['<start_of_text>', '<end_of_text>'] + special_tokens
|
||||
vocab.extend(special_tokens)
|
||||
self.encoder = dict(zip(vocab, range(len(vocab))))
|
||||
self.decoder = {v: k for k, v in self.encoder.items()}
|
||||
self.bpe_ranks = dict(zip(merges, range(len(merges))))
|
||||
self.cache = {t:t for t in special_tokens}
|
||||
special = "|".join(special_tokens)
|
||||
self.pat = re.compile(special + r"""|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
|
||||
|
||||
self.vocab_size = len(self.encoder)
|
||||
self.all_special_ids = [self.encoder[t] for t in special_tokens]
|
||||
|
||||
def bpe(self, token):
|
||||
if token in self.cache:
|
||||
return self.cache[token]
|
||||
word = tuple(token[:-1]) + ( token[-1] + '</w>',)
|
||||
pairs = get_pairs(word)
|
||||
|
||||
if not pairs:
|
||||
return token+'</w>'
|
||||
|
||||
while True:
|
||||
bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
|
||||
if bigram not in self.bpe_ranks:
|
||||
break
|
||||
first, second = bigram
|
||||
new_word = []
|
||||
i = 0
|
||||
while i < len(word):
|
||||
try:
|
||||
j = word.index(first, i)
|
||||
new_word.extend(word[i:j])
|
||||
i = j
|
||||
except:
|
||||
new_word.extend(word[i:])
|
||||
break
|
||||
|
||||
if word[i] == first and i < len(word)-1 and word[i+1] == second:
|
||||
new_word.append(first+second)
|
||||
i += 2
|
||||
else:
|
||||
new_word.append(word[i])
|
||||
i += 1
|
||||
new_word = tuple(new_word)
|
||||
word = new_word
|
||||
if len(word) == 1:
|
||||
break
|
||||
else:
|
||||
pairs = get_pairs(word)
|
||||
word = ' '.join(word)
|
||||
self.cache[token] = word
|
||||
return word
|
||||
|
||||
def encode(self, text):
|
||||
bpe_tokens = []
|
||||
text = whitespace_clean(basic_clean(text)).lower()
|
||||
for token in re.findall(self.pat, text):
|
||||
token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
|
||||
bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
|
||||
return bpe_tokens
|
||||
|
||||
def decode(self, tokens):
|
||||
text = ''.join([self.decoder[token] for token in tokens])
|
||||
text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
|
||||
return text
|
||||
|
||||
|
||||
_tokenizer = SimpleTokenizer()
|
||||
|
||||
|
||||
def tokenize(texts: Union[str, List[str]], context_length: int = 77) -> torch.LongTensor:
|
||||
"""
|
||||
Returns the tokenized representation of given input string(s)
|
||||
|
||||
Parameters
|
||||
----------
|
||||
texts : Union[str, List[str]]
|
||||
An input string or a list of input strings to tokenize
|
||||
context_length : int
|
||||
The context length to use; all CLIP models use 77 as the context length
|
||||
|
||||
Returns
|
||||
-------
|
||||
A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]
|
||||
"""
|
||||
if isinstance(texts, str):
|
||||
texts = [texts]
|
||||
|
||||
sot_token = _tokenizer.encoder["<start_of_text>"]
|
||||
eot_token = _tokenizer.encoder["<end_of_text>"]
|
||||
all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
|
||||
result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
|
||||
|
||||
for i, tokens in enumerate(all_tokens):
|
||||
if len(tokens) > context_length:
|
||||
tokens = tokens[:context_length] # Truncate
|
||||
tokens[-1] = eot_token
|
||||
result[i, :len(tokens)] = torch.tensor(tokens)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class HFTokenizer:
|
||||
"HuggingFace tokenizer wrapper"
|
||||
def __init__(self, tokenizer_name:str):
|
||||
from transformers import AutoTokenizer
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
|
||||
|
||||
def __call__(self, texts:Union[str, List[str]], context_length:int=77) -> torch.Tensor:
|
||||
# same cleaning as for default tokenizer, except lowercasing
|
||||
# adding lower (for case-sensitive tokenizers) will make it more robust but less sensitive to nuance
|
||||
if isinstance(texts, str):
|
||||
texts = [texts]
|
||||
texts = [whitespace_clean(basic_clean(text)) for text in texts]
|
||||
input_ids = self.tokenizer(texts, return_tensors='pt', max_length=context_length, padding='max_length', truncation=True).input_ids
|
||||
return input_ids
|
||||
@@ -0,0 +1,103 @@
|
||||
from typing import Optional, Sequence, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.transforms.functional as F
|
||||
|
||||
from torchvision.transforms import Normalize, Compose, RandomResizedCrop, InterpolationMode, ToTensor, Resize, \
|
||||
CenterCrop
|
||||
|
||||
from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
|
||||
|
||||
|
||||
class ResizeMaxSize(nn.Module):
|
||||
|
||||
def __init__(self, max_size, interpolation=InterpolationMode.BICUBIC, fn='max', fill=0):
|
||||
super().__init__()
|
||||
if not isinstance(max_size, int):
|
||||
raise TypeError(f"Size should be int. Got {type(max_size)}")
|
||||
self.max_size = max_size
|
||||
self.interpolation = interpolation
|
||||
self.fn = min if fn == 'min' else min
|
||||
self.fill = fill
|
||||
|
||||
def forward(self, img):
|
||||
if isinstance(img, torch.Tensor):
|
||||
height, width = img.shape[:2]
|
||||
else:
|
||||
width, height = img.size
|
||||
scale = self.max_size / float(max(height, width))
|
||||
if scale != 1.0:
|
||||
new_size = tuple(round(dim * scale) for dim in (height, width))
|
||||
img = F.resize(img, new_size, self.interpolation)
|
||||
pad_h = self.max_size - new_size[0]
|
||||
pad_w = self.max_size - new_size[1]
|
||||
img = F.pad(img, padding=[pad_w//2, pad_h//2, pad_w - pad_w//2, pad_h - pad_h//2], fill=self.fill)
|
||||
return img
|
||||
|
||||
|
||||
def _convert_to_rgb(image):
|
||||
return image.convert('RGB')
|
||||
|
||||
|
||||
# class CatGen(nn.Module):
|
||||
# def __init__(self, num=4):
|
||||
# self.num = num
|
||||
# def mixgen_batch(image, text):
|
||||
# batch_size = image.shape[0]
|
||||
# index = np.random.permutation(batch_size)
|
||||
|
||||
# cat_images = []
|
||||
# for i in range(batch_size):
|
||||
# # image mixup
|
||||
# image[i,:] = lam * image[i,:] + (1 - lam) * image[index[i],:]
|
||||
# # text concat
|
||||
# text[i] = tokenizer((str(text[i]) + " " + str(text[index[i]])))[0]
|
||||
# text = torch.stack(text)
|
||||
# return image, text
|
||||
|
||||
|
||||
def image_transform(
|
||||
image_size: int,
|
||||
is_train: bool,
|
||||
mean: Optional[Tuple[float, ...]] = None,
|
||||
std: Optional[Tuple[float, ...]] = None,
|
||||
resize_longest_max: bool = False,
|
||||
fill_color: int = 0,
|
||||
):
|
||||
mean = mean or OPENAI_DATASET_MEAN
|
||||
if not isinstance(mean, (list, tuple)):
|
||||
mean = (mean,) * 3
|
||||
|
||||
std = std or OPENAI_DATASET_STD
|
||||
if not isinstance(std, (list, tuple)):
|
||||
std = (std,) * 3
|
||||
|
||||
if isinstance(image_size, (list, tuple)) and image_size[0] == image_size[1]:
|
||||
# for square size, pass size as int so that Resize() uses aspect preserving shortest edge
|
||||
image_size = image_size[0]
|
||||
|
||||
normalize = Normalize(mean=mean, std=std)
|
||||
if is_train:
|
||||
return Compose([
|
||||
RandomResizedCrop(image_size, scale=(0.9, 1.0), interpolation=InterpolationMode.BICUBIC),
|
||||
_convert_to_rgb,
|
||||
ToTensor(),
|
||||
normalize,
|
||||
])
|
||||
else:
|
||||
if resize_longest_max:
|
||||
transforms = [
|
||||
ResizeMaxSize(image_size, fill=fill_color)
|
||||
]
|
||||
else:
|
||||
transforms = [
|
||||
Resize(image_size, interpolation=InterpolationMode.BICUBIC),
|
||||
CenterCrop(image_size),
|
||||
]
|
||||
transforms.extend([
|
||||
_convert_to_rgb,
|
||||
ToTensor(),
|
||||
normalize,
|
||||
])
|
||||
return Compose(transforms)
|
||||
@@ -0,0 +1,737 @@
|
||||
import os
|
||||
import logging
|
||||
from collections import OrderedDict
|
||||
import math
|
||||
from typing import Callable, Optional, Sequence
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
try:
|
||||
from timm.models.layers import trunc_normal_
|
||||
except:
|
||||
from timm.layers import trunc_normal_
|
||||
|
||||
from .rope import VisionRotaryEmbedding, VisionRotaryEmbeddingFast
|
||||
from .utils import to_2tuple
|
||||
|
||||
if os.getenv('ENV_TYPE') == 'deepspeed':
|
||||
try:
|
||||
import deepspeed
|
||||
from deepspeed.runtime.activation_checkpointing.checkpointing import checkpoint
|
||||
except:
|
||||
print("Please 'pip install deepspeed'")
|
||||
deepspeed = None
|
||||
from torch.utils.checkpoint import checkpoint
|
||||
else:
|
||||
from torch.utils.checkpoint import checkpoint
|
||||
|
||||
try:
|
||||
import xformers.ops as xops
|
||||
except ImportError:
|
||||
xops = None
|
||||
# print("Please 'pip install xformers'")
|
||||
|
||||
class LayerNormFp32(nn.LayerNorm):
|
||||
"""Subclass torch's LayerNorm to handle fp16 (by casting to float32 and back)."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
output = F.layer_norm(
|
||||
x.float(),
|
||||
self.normalized_shape,
|
||||
self.weight.float() if self.weight is not None else None,
|
||||
self.bias.float() if self.bias is not None else None,
|
||||
self.eps,
|
||||
)
|
||||
return output.type_as(x)
|
||||
|
||||
|
||||
class LayerNorm(nn.LayerNorm):
|
||||
"""Subclass torch's LayerNorm (with cast back to input dtype)."""
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
orig_type = x.dtype
|
||||
x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
|
||||
return x.to(orig_type)
|
||||
|
||||
class QuickGELU(nn.Module):
|
||||
# NOTE This is slower than nn.GELU or nn.SiLU and uses more GPU memory
|
||||
def forward(self, x: torch.Tensor):
|
||||
return x * torch.sigmoid(1.702 * x)
|
||||
|
||||
|
||||
class LayerScale(nn.Module):
|
||||
def __init__(self, dim, init_values=1e-5, inplace=False):
|
||||
super().__init__()
|
||||
self.inplace = inplace
|
||||
self.gamma = nn.Parameter(init_values * torch.ones(dim))
|
||||
|
||||
def forward(self, x):
|
||||
return x.mul_(self.gamma) if self.inplace else x * self.gamma
|
||||
|
||||
class PatchDropout(nn.Module):
|
||||
"""
|
||||
https://arxiv.org/abs/2212.00794
|
||||
"""
|
||||
|
||||
def __init__(self, prob, exclude_first_token=True):
|
||||
super().__init__()
|
||||
assert 0 <= prob < 1.
|
||||
self.prob = prob
|
||||
self.exclude_first_token = exclude_first_token # exclude CLS token
|
||||
logging.info(f"os.getenv('RoPE')={os.getenv('RoPE')}")
|
||||
|
||||
def forward(self, x):
|
||||
if not self.training or self.prob == 0.:
|
||||
return x
|
||||
|
||||
if self.exclude_first_token:
|
||||
cls_tokens, x = x[:, :1], x[:, 1:]
|
||||
else:
|
||||
cls_tokens = torch.jit.annotate(torch.Tensor, x[:, :1])
|
||||
|
||||
batch = x.size()[0]
|
||||
num_tokens = x.size()[1]
|
||||
|
||||
batch_indices = torch.arange(batch)
|
||||
batch_indices = batch_indices[..., None]
|
||||
|
||||
keep_prob = 1 - self.prob
|
||||
num_patches_keep = max(1, int(num_tokens * keep_prob))
|
||||
|
||||
rand = torch.randn(batch, num_tokens)
|
||||
patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices
|
||||
|
||||
x = x[batch_indices, patch_indices_keep]
|
||||
|
||||
if self.exclude_first_token:
|
||||
x = torch.cat((cls_tokens, x), dim=1)
|
||||
|
||||
if self.training and os.getenv('RoPE') == '1':
|
||||
return x, patch_indices_keep
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def _in_projection_packed(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
w: torch.Tensor,
|
||||
b: Optional[torch.Tensor] = None,
|
||||
):
|
||||
"""
|
||||
https://github.com/pytorch/pytorch/blob/db2a237763eb8693a20788be94f8c192e762baa8/torch/nn/functional.py#L4726
|
||||
"""
|
||||
E = q.size(-1)
|
||||
if k is v:
|
||||
if q is k:
|
||||
# self-attention
|
||||
return F.linear(q, w, b).chunk(3, dim=-1)
|
||||
else:
|
||||
# encoder-decoder attention
|
||||
w_q, w_kv = w.split([E, E * 2])
|
||||
if b is None:
|
||||
b_q = b_kv = None
|
||||
else:
|
||||
b_q, b_kv = b.split([E, E * 2])
|
||||
return (F.linear(q, w_q, b_q),) + F.linear(k, w_kv, b_kv).chunk(2, dim=-1)
|
||||
else:
|
||||
w_q, w_k, w_v = w.chunk(3)
|
||||
if b is None:
|
||||
b_q = b_k = b_v = None
|
||||
else:
|
||||
b_q, b_k, b_v = b.chunk(3)
|
||||
return F.linear(q, w_q, b_q), F.linear(k, w_k, b_k), F.linear(v, w_v, b_v)
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
num_heads=8,
|
||||
qkv_bias=True,
|
||||
scaled_cosine=False,
|
||||
scale_heads=False,
|
||||
logit_scale_max=math.log(1. / 0.01),
|
||||
attn_drop=0.,
|
||||
proj_drop=0.,
|
||||
xattn=False,
|
||||
rope=False
|
||||
):
|
||||
super().__init__()
|
||||
self.scaled_cosine = scaled_cosine
|
||||
self.scale_heads = scale_heads
|
||||
assert dim % num_heads == 0, 'dim should be divisible by num_heads'
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = dim // num_heads
|
||||
self.scale = self.head_dim ** -0.5
|
||||
self.logit_scale_max = logit_scale_max
|
||||
|
||||
# keeping in_proj in this form (instead of nn.Linear) to match weight scheme of original
|
||||
self.in_proj_weight = nn.Parameter(torch.randn((dim * 3, dim)) * self.scale)
|
||||
if qkv_bias:
|
||||
self.in_proj_bias = nn.Parameter(torch.zeros(dim * 3))
|
||||
else:
|
||||
self.in_proj_bias = None
|
||||
|
||||
if self.scaled_cosine:
|
||||
self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))))
|
||||
else:
|
||||
self.logit_scale = None
|
||||
self.attn_drop = nn.Dropout(attn_drop)
|
||||
if self.scale_heads:
|
||||
self.head_scale = nn.Parameter(torch.ones((num_heads, 1, 1)))
|
||||
else:
|
||||
self.head_scale = None
|
||||
self.out_proj = nn.Linear(dim, dim)
|
||||
self.out_drop = nn.Dropout(proj_drop)
|
||||
self.xattn = xattn
|
||||
self.xattn_drop = attn_drop
|
||||
self.rope = rope
|
||||
|
||||
def forward(self, x, attn_mask: Optional[torch.Tensor] = None):
|
||||
L, N, C = x.shape
|
||||
q, k, v = F.linear(x, self.in_proj_weight, self.in_proj_bias).chunk(3, dim=-1)
|
||||
if self.xattn:
|
||||
q = q.contiguous().view(L, N, self.num_heads, -1).transpose(0, 1)
|
||||
k = k.contiguous().view(L, N, self.num_heads, -1).transpose(0, 1)
|
||||
v = v.contiguous().view(L, N, self.num_heads, -1).transpose(0, 1)
|
||||
|
||||
x = xops.memory_efficient_attention(
|
||||
q, k, v,
|
||||
p=self.xattn_drop,
|
||||
scale=self.scale if self.logit_scale is None else None,
|
||||
attn_bias=xops.LowerTriangularMask() if attn_mask is not None else None,
|
||||
)
|
||||
else:
|
||||
q = q.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
|
||||
k = k.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
|
||||
v = v.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
|
||||
|
||||
if self.logit_scale is not None:
|
||||
attn = torch.bmm(F.normalize(q, dim=-1), F.normalize(k, dim=-1).transpose(-1, -2))
|
||||
logit_scale = torch.clamp(self.logit_scale, max=self.logit_scale_max).exp()
|
||||
attn = attn.view(N, self.num_heads, L, L) * logit_scale
|
||||
attn = attn.view(-1, L, L)
|
||||
else:
|
||||
q = q * self.scale
|
||||
attn = torch.bmm(q, k.transpose(-1, -2))
|
||||
|
||||
if attn_mask is not None:
|
||||
if attn_mask.dtype == torch.bool:
|
||||
new_attn_mask = torch.zeros_like(attn_mask, dtype=q.dtype)
|
||||
new_attn_mask.masked_fill_(attn_mask, float("-inf"))
|
||||
attn_mask = new_attn_mask
|
||||
attn += attn_mask
|
||||
|
||||
attn = attn.softmax(dim=-1)
|
||||
attn = self.attn_drop(attn)
|
||||
|
||||
x = torch.bmm(attn, v)
|
||||
|
||||
if self.head_scale is not None:
|
||||
x = x.view(N, self.num_heads, L, C) * self.head_scale
|
||||
x = x.view(-1, L, C)
|
||||
x = x.transpose(0, 1).reshape(L, N, C)
|
||||
x = self.out_proj(x)
|
||||
x = self.out_drop(x)
|
||||
return x
|
||||
|
||||
class CustomAttention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
num_heads=8,
|
||||
qkv_bias=True,
|
||||
scaled_cosine=True,
|
||||
scale_heads=False,
|
||||
logit_scale_max=math.log(1. / 0.01),
|
||||
attn_drop=0.,
|
||||
proj_drop=0.,
|
||||
xattn=False
|
||||
):
|
||||
super().__init__()
|
||||
self.scaled_cosine = scaled_cosine
|
||||
self.scale_heads = scale_heads
|
||||
assert dim % num_heads == 0, 'dim should be divisible by num_heads'
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = dim // num_heads
|
||||
self.scale = self.head_dim ** -0.5
|
||||
self.logit_scale_max = logit_scale_max
|
||||
|
||||
# keeping in_proj in this form (instead of nn.Linear) to match weight scheme of original
|
||||
self.in_proj_weight = nn.Parameter(torch.randn((dim * 3, dim)) * self.scale)
|
||||
if qkv_bias:
|
||||
self.in_proj_bias = nn.Parameter(torch.zeros(dim * 3))
|
||||
else:
|
||||
self.in_proj_bias = None
|
||||
|
||||
if self.scaled_cosine:
|
||||
self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))))
|
||||
else:
|
||||
self.logit_scale = None
|
||||
self.attn_drop = nn.Dropout(attn_drop)
|
||||
if self.scale_heads:
|
||||
self.head_scale = nn.Parameter(torch.ones((num_heads, 1, 1)))
|
||||
else:
|
||||
self.head_scale = None
|
||||
self.out_proj = nn.Linear(dim, dim)
|
||||
self.out_drop = nn.Dropout(proj_drop)
|
||||
self.xattn = xattn
|
||||
self.xattn_drop = attn_drop
|
||||
|
||||
def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
|
||||
q, k, v = _in_projection_packed(query, key, value, self.in_proj_weight, self.in_proj_bias)
|
||||
N_q, B_q, C_q = q.shape
|
||||
N_k, B_k, C_k = k.shape
|
||||
N_v, B_v, C_v = v.shape
|
||||
if self.xattn:
|
||||
# B, N, C -> B, N, num_heads, C
|
||||
q = q.permute(1, 0, 2).reshape(B_q, N_q, self.num_heads, -1)
|
||||
k = k.permute(1, 0, 2).reshape(B_k, N_k, self.num_heads, -1)
|
||||
v = v.permute(1, 0, 2).reshape(B_v, N_v, self.num_heads, -1)
|
||||
|
||||
x = xops.memory_efficient_attention(
|
||||
q, k, v,
|
||||
p=self.xattn_drop,
|
||||
scale=self.scale if self.logit_scale is None else None,
|
||||
attn_bias=xops.LowerTriangularMask() if attn_mask is not None else None
|
||||
)
|
||||
else:
|
||||
# B*H, L, C
|
||||
q = q.contiguous().view(N_q, B_q * self.num_heads, -1).transpose(0, 1)
|
||||
k = k.contiguous().view(N_k, B_k * self.num_heads, -1).transpose(0, 1)
|
||||
v = v.contiguous().view(N_v, B_v * self.num_heads, -1).transpose(0, 1)
|
||||
|
||||
if self.logit_scale is not None:
|
||||
# B*H, N_q, N_k
|
||||
attn = torch.bmm(F.normalize(q, dim=-1), F.normalize(k, dim=-1).transpose(-1, -2))
|
||||
logit_scale = torch.clamp(self.logit_scale, max=self.logit_scale_max).exp()
|
||||
attn = attn.view(B_q, self.num_heads, N_q, N_k) * logit_scale
|
||||
attn = attn.view(-1, N_q, N_k)
|
||||
else:
|
||||
q = q * self.scale
|
||||
attn = torch.bmm(q, k.transpose(-1, -2))
|
||||
|
||||
if attn_mask is not None:
|
||||
if attn_mask.dtype == torch.bool:
|
||||
new_attn_mask = torch.zeros_like(attn_mask, dtype=q.dtype)
|
||||
new_attn_mask.masked_fill_(attn_mask, float("-inf"))
|
||||
attn_mask = new_attn_mask
|
||||
attn += attn_mask
|
||||
|
||||
attn = attn.softmax(dim=-1)
|
||||
attn = self.attn_drop(attn)
|
||||
|
||||
x = torch.bmm(attn, v)
|
||||
|
||||
if self.head_scale is not None:
|
||||
x = x.view(B_q, self.num_heads, N_q, C_q) * self.head_scale
|
||||
x = x.view(-1, N_q, C_q)
|
||||
x = x.transpose(0, 1).reshape(N_q, B_q, C_q)
|
||||
x = self.out_proj(x)
|
||||
x = self.out_drop(x)
|
||||
return x
|
||||
|
||||
class CustomResidualAttentionBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int,
|
||||
n_head: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
ls_init_value: float = None,
|
||||
act_layer: Callable = nn.GELU,
|
||||
norm_layer: Callable = LayerNorm,
|
||||
scale_cosine_attn: bool = False,
|
||||
scale_heads: bool = False,
|
||||
scale_attn: bool = False,
|
||||
scale_fc: bool = False,
|
||||
cross_attn: bool = False,
|
||||
xattn: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.ln_1 = norm_layer(d_model)
|
||||
self.ln_1_k = norm_layer(d_model) if cross_attn else self.ln_1
|
||||
self.ln_1_v = norm_layer(d_model) if cross_attn else self.ln_1
|
||||
self.attn = CustomAttention(
|
||||
d_model, n_head,
|
||||
qkv_bias=True,
|
||||
attn_drop=0.,
|
||||
proj_drop=0.,
|
||||
scaled_cosine=scale_cosine_attn,
|
||||
scale_heads=scale_heads,
|
||||
xattn=xattn
|
||||
)
|
||||
|
||||
self.ln_attn = norm_layer(d_model) if scale_attn else nn.Identity()
|
||||
self.ls_1 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
|
||||
|
||||
self.ln_2 = norm_layer(d_model)
|
||||
mlp_width = int(d_model * mlp_ratio)
|
||||
self.mlp = nn.Sequential(OrderedDict([
|
||||
("c_fc", nn.Linear(d_model, mlp_width)),
|
||||
('ln', norm_layer(mlp_width) if scale_fc else nn.Identity()),
|
||||
("gelu", act_layer()),
|
||||
("c_proj", nn.Linear(mlp_width, d_model))
|
||||
]))
|
||||
|
||||
self.ls_2 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
|
||||
|
||||
def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
|
||||
q = q + self.ls_1(self.ln_attn(self.attn(self.ln_1(q), self.ln_1_k(k), self.ln_1_v(v), attn_mask=attn_mask)))
|
||||
q = q + self.ls_2(self.mlp(self.ln_2(q)))
|
||||
return q
|
||||
|
||||
class CustomTransformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
width: int,
|
||||
layers: int,
|
||||
heads: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
ls_init_value: float = None,
|
||||
act_layer: Callable = nn.GELU,
|
||||
norm_layer: Callable = LayerNorm,
|
||||
scale_cosine_attn: bool = True,
|
||||
scale_heads: bool = False,
|
||||
scale_attn: bool = False,
|
||||
scale_fc: bool = False,
|
||||
cross_attn: bool = False,
|
||||
xattn: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.width = width
|
||||
self.layers = layers
|
||||
self.grad_checkpointing = False
|
||||
self.xattn = xattn
|
||||
|
||||
self.resblocks = nn.ModuleList([
|
||||
CustomResidualAttentionBlock(
|
||||
width,
|
||||
heads,
|
||||
mlp_ratio,
|
||||
ls_init_value=ls_init_value,
|
||||
act_layer=act_layer,
|
||||
norm_layer=norm_layer,
|
||||
scale_cosine_attn=scale_cosine_attn,
|
||||
scale_heads=scale_heads,
|
||||
scale_attn=scale_attn,
|
||||
scale_fc=scale_fc,
|
||||
cross_attn=cross_attn,
|
||||
xattn=xattn)
|
||||
for _ in range(layers)
|
||||
])
|
||||
|
||||
def get_cast_dtype(self) -> torch.dtype:
|
||||
return self.resblocks[0].mlp.c_fc.weight.dtype
|
||||
|
||||
def forward(self, q: torch.Tensor, k: torch.Tensor = None, v: torch.Tensor = None, attn_mask: Optional[torch.Tensor] = None):
|
||||
if k is None and v is None:
|
||||
k = v = q
|
||||
for r in self.resblocks:
|
||||
if self.grad_checkpointing and not torch.jit.is_scripting():
|
||||
q = checkpoint(r, q, k, v, attn_mask)
|
||||
else:
|
||||
q = r(q, k, v, attn_mask=attn_mask)
|
||||
return q
|
||||
|
||||
|
||||
class ResidualAttentionBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int,
|
||||
n_head: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
ls_init_value: float = None,
|
||||
act_layer: Callable = nn.GELU,
|
||||
norm_layer: Callable = LayerNorm,
|
||||
xattn: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.ln_1 = norm_layer(d_model)
|
||||
if xattn:
|
||||
self.attn = Attention(d_model, n_head, xattn=True)
|
||||
else:
|
||||
self.attn = nn.MultiheadAttention(d_model, n_head)
|
||||
self.ls_1 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
|
||||
|
||||
self.ln_2 = norm_layer(d_model)
|
||||
mlp_width = int(d_model * mlp_ratio)
|
||||
self.mlp = nn.Sequential(OrderedDict([
|
||||
("c_fc", nn.Linear(d_model, mlp_width)),
|
||||
("gelu", act_layer()),
|
||||
("c_proj", nn.Linear(mlp_width, d_model))
|
||||
]))
|
||||
|
||||
self.ls_2 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
|
||||
self.xattn = xattn
|
||||
|
||||
def attention(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
|
||||
attn_mask = attn_mask.to(x.dtype) if attn_mask is not None else None
|
||||
if self.xattn:
|
||||
return self.attn(x, attn_mask=attn_mask)
|
||||
return self.attn(x, x, x, need_weights=False, attn_mask=attn_mask)[0]
|
||||
|
||||
def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
|
||||
x = x + self.ls_1(self.attention(self.ln_1(x), attn_mask=attn_mask))
|
||||
x = x + self.ls_2(self.mlp(self.ln_2(x)))
|
||||
return x
|
||||
|
||||
class Transformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
width: int,
|
||||
layers: int,
|
||||
heads: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
ls_init_value: float = None,
|
||||
act_layer: Callable = nn.GELU,
|
||||
norm_layer: Callable = LayerNorm,
|
||||
xattn: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.width = width
|
||||
self.layers = layers
|
||||
self.grad_checkpointing = False
|
||||
|
||||
self.resblocks = nn.ModuleList([
|
||||
ResidualAttentionBlock(
|
||||
width, heads, mlp_ratio, ls_init_value=ls_init_value, act_layer=act_layer, norm_layer=norm_layer, xattn=xattn)
|
||||
for _ in range(layers)
|
||||
])
|
||||
|
||||
def get_cast_dtype(self) -> torch.dtype:
|
||||
return self.resblocks[0].mlp.c_fc.weight.dtype
|
||||
|
||||
def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
|
||||
for r in self.resblocks:
|
||||
if self.grad_checkpointing and not torch.jit.is_scripting():
|
||||
x = checkpoint(r, x, attn_mask)
|
||||
else:
|
||||
x = r(x, attn_mask=attn_mask)
|
||||
return x
|
||||
|
||||
|
||||
class VisionTransformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
image_size: int,
|
||||
patch_size: int,
|
||||
width: int,
|
||||
layers: int,
|
||||
heads: int,
|
||||
mlp_ratio: float,
|
||||
ls_init_value: float = None,
|
||||
patch_dropout: float = 0.,
|
||||
global_average_pool: bool = False,
|
||||
output_dim: int = 512,
|
||||
act_layer: Callable = nn.GELU,
|
||||
norm_layer: Callable = LayerNorm,
|
||||
xattn: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.image_size = to_2tuple(image_size)
|
||||
self.patch_size = to_2tuple(patch_size)
|
||||
self.grid_size = (self.image_size[0] // self.patch_size[0], self.image_size[1] // self.patch_size[1])
|
||||
self.output_dim = output_dim
|
||||
self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
|
||||
|
||||
scale = width ** -0.5
|
||||
self.class_embedding = nn.Parameter(scale * torch.randn(width))
|
||||
self.positional_embedding = nn.Parameter(scale * torch.randn(self.grid_size[0] * self.grid_size[1] + 1, width))
|
||||
|
||||
# setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
|
||||
self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()
|
||||
self.ln_pre = norm_layer(width)
|
||||
|
||||
self.transformer = Transformer(
|
||||
width,
|
||||
layers,
|
||||
heads,
|
||||
mlp_ratio,
|
||||
ls_init_value=ls_init_value,
|
||||
act_layer=act_layer,
|
||||
norm_layer=norm_layer,
|
||||
xattn=xattn
|
||||
)
|
||||
|
||||
self.global_average_pool = global_average_pool
|
||||
self.ln_post = norm_layer(width)
|
||||
self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
|
||||
|
||||
def lock(self, unlocked_groups=0, freeze_bn_stats=False):
|
||||
for param in self.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
if unlocked_groups != 0:
|
||||
groups = [
|
||||
[
|
||||
self.conv1,
|
||||
self.class_embedding,
|
||||
self.positional_embedding,
|
||||
self.ln_pre,
|
||||
],
|
||||
*self.transformer.resblocks[:-1],
|
||||
[
|
||||
self.transformer.resblocks[-1],
|
||||
self.ln_post,
|
||||
],
|
||||
self.proj,
|
||||
]
|
||||
|
||||
def _unlock(x):
|
||||
if isinstance(x, Sequence):
|
||||
for g in x:
|
||||
_unlock(g)
|
||||
else:
|
||||
if isinstance(x, torch.nn.Parameter):
|
||||
x.requires_grad = True
|
||||
else:
|
||||
for p in x.parameters():
|
||||
p.requires_grad = True
|
||||
|
||||
_unlock(groups[-unlocked_groups:])
|
||||
|
||||
def get_num_layers(self):
|
||||
return self.transformer.layers
|
||||
|
||||
@torch.jit.ignore
|
||||
def set_grad_checkpointing(self, enable=True):
|
||||
self.transformer.grad_checkpointing = enable
|
||||
|
||||
@torch.jit.ignore
|
||||
def no_weight_decay(self):
|
||||
return {'positional_embedding', 'class_embedding'}
|
||||
|
||||
def forward(self, x: torch.Tensor, return_all_features: bool=False):
|
||||
x = self.conv1(x) # shape = [*, width, grid, grid]
|
||||
x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [*, width, grid ** 2]
|
||||
x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width]
|
||||
x = torch.cat(
|
||||
[self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device),
|
||||
x], dim=1) # shape = [*, grid ** 2 + 1, width]
|
||||
x = x + self.positional_embedding.to(x.dtype)
|
||||
|
||||
# a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
|
||||
x = self.patch_dropout(x)
|
||||
x = self.ln_pre(x)
|
||||
|
||||
x = x.permute(1, 0, 2) # NLD -> LND
|
||||
x = self.transformer(x)
|
||||
x = x.permute(1, 0, 2) # LND -> NLD
|
||||
|
||||
if not return_all_features:
|
||||
if self.global_average_pool:
|
||||
x = x.mean(dim=1) #x = x[:,1:,:].mean(dim=1)
|
||||
else:
|
||||
x = x[:, 0]
|
||||
|
||||
x = self.ln_post(x)
|
||||
|
||||
if self.proj is not None:
|
||||
x = x @ self.proj
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class TextTransformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
context_length: int = 77,
|
||||
vocab_size: int = 49408,
|
||||
width: int = 512,
|
||||
heads: int = 8,
|
||||
layers: int = 12,
|
||||
ls_init_value: float = None,
|
||||
output_dim: int = 512,
|
||||
act_layer: Callable = nn.GELU,
|
||||
norm_layer: Callable = LayerNorm,
|
||||
xattn: bool= False,
|
||||
attn_mask: bool = True
|
||||
):
|
||||
super().__init__()
|
||||
self.context_length = context_length
|
||||
self.vocab_size = vocab_size
|
||||
self.width = width
|
||||
self.output_dim = output_dim
|
||||
|
||||
self.token_embedding = nn.Embedding(vocab_size, width)
|
||||
self.positional_embedding = nn.Parameter(torch.empty(self.context_length, width))
|
||||
self.transformer = Transformer(
|
||||
width=width,
|
||||
layers=layers,
|
||||
heads=heads,
|
||||
ls_init_value=ls_init_value,
|
||||
act_layer=act_layer,
|
||||
norm_layer=norm_layer,
|
||||
xattn=xattn
|
||||
)
|
||||
|
||||
self.xattn = xattn
|
||||
self.ln_final = norm_layer(width)
|
||||
self.text_projection = nn.Parameter(torch.empty(width, output_dim))
|
||||
|
||||
if attn_mask:
|
||||
self.register_buffer('attn_mask', self.build_attention_mask(), persistent=False)
|
||||
else:
|
||||
self.attn_mask = None
|
||||
|
||||
self.init_parameters()
|
||||
|
||||
def init_parameters(self):
|
||||
nn.init.normal_(self.token_embedding.weight, std=0.02)
|
||||
nn.init.normal_(self.positional_embedding, std=0.01)
|
||||
|
||||
proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
|
||||
attn_std = self.transformer.width ** -0.5
|
||||
fc_std = (2 * self.transformer.width) ** -0.5
|
||||
for block in self.transformer.resblocks:
|
||||
nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
|
||||
nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
|
||||
nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
|
||||
nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
|
||||
|
||||
if self.text_projection is not None:
|
||||
nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
|
||||
|
||||
@torch.jit.ignore
|
||||
def set_grad_checkpointing(self, enable=True):
|
||||
self.transformer.grad_checkpointing = enable
|
||||
|
||||
@torch.jit.ignore
|
||||
def no_weight_decay(self):
|
||||
# return {'positional_embedding', 'token_embedding'}
|
||||
return {'positional_embedding'}
|
||||
|
||||
def get_num_layers(self):
|
||||
return self.transformer.layers
|
||||
|
||||
def build_attention_mask(self):
|
||||
# lazily create causal attention mask, with full attention between the vision tokens
|
||||
# pytorch uses additive attention mask; fill with -inf
|
||||
mask = torch.empty(self.context_length, self.context_length)
|
||||
mask.fill_(float("-inf"))
|
||||
mask.triu_(1) # zero out the lower diagonal
|
||||
return mask
|
||||
|
||||
def forward(self, text, return_all_features: bool=False):
|
||||
cast_dtype = self.transformer.get_cast_dtype()
|
||||
x = self.token_embedding(text).to(cast_dtype) # [batch_size, n_ctx, d_model]
|
||||
|
||||
x = x + self.positional_embedding.to(cast_dtype)
|
||||
x = x.permute(1, 0, 2) # NLD -> LND
|
||||
x = self.transformer(x, attn_mask=self.attn_mask)
|
||||
# x = self.transformer(x) # no attention mask is applied
|
||||
x = x.permute(1, 0, 2) # LND -> NLD
|
||||
x = self.ln_final(x)
|
||||
|
||||
if not return_all_features:
|
||||
# x.shape = [batch_size, n_ctx, transformer.width]
|
||||
# take features from the eot embedding (eot_token is the highest number in each sequence)
|
||||
x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
|
||||
return x
|
||||
@@ -0,0 +1,326 @@
|
||||
from itertools import repeat
|
||||
import collections.abc
|
||||
import logging
|
||||
import math
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
from torch import nn as nn
|
||||
from torchvision.ops.misc import FrozenBatchNorm2d
|
||||
import torch.nn.functional as F
|
||||
|
||||
# open CLIP
|
||||
def resize_clip_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_dim=1):
|
||||
# Rescale the grid of position embeddings when loading from state_dict
|
||||
old_pos_embed = state_dict.get('visual.positional_embedding', None)
|
||||
if old_pos_embed is None or not hasattr(model.visual, 'grid_size'):
|
||||
return
|
||||
grid_size = to_2tuple(model.visual.grid_size)
|
||||
extra_tokens = 1 # FIXME detect different token configs (ie no class token, or more)
|
||||
new_seq_len = grid_size[0] * grid_size[1] + extra_tokens
|
||||
if new_seq_len == old_pos_embed.shape[0]:
|
||||
return
|
||||
|
||||
if extra_tokens:
|
||||
pos_emb_tok, pos_emb_img = old_pos_embed[:extra_tokens], old_pos_embed[extra_tokens:]
|
||||
else:
|
||||
pos_emb_tok, pos_emb_img = None, old_pos_embed
|
||||
old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img))))
|
||||
|
||||
logging.info('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size)
|
||||
pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
|
||||
pos_emb_img = F.interpolate(
|
||||
pos_emb_img,
|
||||
size=grid_size,
|
||||
mode=interpolation,
|
||||
align_corners=True,
|
||||
)
|
||||
pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)[0]
|
||||
if pos_emb_tok is not None:
|
||||
new_pos_embed = torch.cat([pos_emb_tok, pos_emb_img], dim=0)
|
||||
else:
|
||||
new_pos_embed = pos_emb_img
|
||||
state_dict['visual.positional_embedding'] = new_pos_embed
|
||||
|
||||
|
||||
def resize_visual_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_dim=1):
|
||||
# Rescale the grid of position embeddings when loading from state_dict
|
||||
old_pos_embed = state_dict.get('positional_embedding', None)
|
||||
if old_pos_embed is None or not hasattr(model.visual, 'grid_size'):
|
||||
return
|
||||
grid_size = to_2tuple(model.visual.grid_size)
|
||||
extra_tokens = 1 # FIXME detect different token configs (ie no class token, or more)
|
||||
new_seq_len = grid_size[0] * grid_size[1] + extra_tokens
|
||||
if new_seq_len == old_pos_embed.shape[0]:
|
||||
return
|
||||
|
||||
if extra_tokens:
|
||||
pos_emb_tok, pos_emb_img = old_pos_embed[:extra_tokens], old_pos_embed[extra_tokens:]
|
||||
else:
|
||||
pos_emb_tok, pos_emb_img = None, old_pos_embed
|
||||
old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img))))
|
||||
|
||||
logging.info('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size)
|
||||
pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
|
||||
pos_emb_img = F.interpolate(
|
||||
pos_emb_img,
|
||||
size=grid_size,
|
||||
mode=interpolation,
|
||||
align_corners=True,
|
||||
)
|
||||
pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)[0]
|
||||
if pos_emb_tok is not None:
|
||||
new_pos_embed = torch.cat([pos_emb_tok, pos_emb_img], dim=0)
|
||||
else:
|
||||
new_pos_embed = pos_emb_img
|
||||
state_dict['positional_embedding'] = new_pos_embed
|
||||
|
||||
def resize_evaclip_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_dim=1):
|
||||
all_keys = list(state_dict.keys())
|
||||
# interpolate position embedding
|
||||
if 'visual.pos_embed' in state_dict:
|
||||
pos_embed_checkpoint = state_dict['visual.pos_embed']
|
||||
embedding_size = pos_embed_checkpoint.shape[-1]
|
||||
num_patches = model.visual.patch_embed.num_patches
|
||||
num_extra_tokens = model.visual.pos_embed.shape[-2] - num_patches
|
||||
# height (== width) for the checkpoint position embedding
|
||||
orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
|
||||
# height (== width) for the new position embedding
|
||||
new_size = int(num_patches ** 0.5)
|
||||
# class_token and dist_token are kept unchanged
|
||||
if orig_size != new_size:
|
||||
print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
|
||||
extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
|
||||
# only the position tokens are interpolated
|
||||
pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
|
||||
pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
|
||||
pos_tokens = torch.nn.functional.interpolate(
|
||||
pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
|
||||
pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
|
||||
new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
|
||||
state_dict['visual.pos_embed'] = new_pos_embed
|
||||
|
||||
patch_embed_proj = state_dict['visual.patch_embed.proj.weight']
|
||||
patch_size = model.visual.patch_embed.patch_size
|
||||
state_dict['visual.patch_embed.proj.weight'] = torch.nn.functional.interpolate(
|
||||
patch_embed_proj.float(), size=patch_size, mode='bicubic', align_corners=False)
|
||||
|
||||
|
||||
def resize_eva_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_dim=1):
|
||||
all_keys = list(state_dict.keys())
|
||||
# interpolate position embedding
|
||||
if 'pos_embed' in state_dict:
|
||||
pos_embed_checkpoint = state_dict['pos_embed']
|
||||
embedding_size = pos_embed_checkpoint.shape[-1]
|
||||
num_patches = model.visual.patch_embed.num_patches
|
||||
num_extra_tokens = model.visual.pos_embed.shape[-2] - num_patches
|
||||
# height (== width) for the checkpoint position embedding
|
||||
orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
|
||||
# height (== width) for the new position embedding
|
||||
new_size = int(num_patches ** 0.5)
|
||||
# class_token and dist_token are kept unchanged
|
||||
if orig_size != new_size:
|
||||
print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
|
||||
extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
|
||||
# only the position tokens are interpolated
|
||||
pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
|
||||
pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
|
||||
pos_tokens = torch.nn.functional.interpolate(
|
||||
pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
|
||||
pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
|
||||
new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
|
||||
state_dict['pos_embed'] = new_pos_embed
|
||||
|
||||
patch_embed_proj = state_dict['patch_embed.proj.weight']
|
||||
patch_size = model.visual.patch_embed.patch_size
|
||||
state_dict['patch_embed.proj.weight'] = torch.nn.functional.interpolate(
|
||||
patch_embed_proj.float(), size=patch_size, mode='bicubic', align_corners=False)
|
||||
|
||||
|
||||
def resize_rel_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_dim=1):
|
||||
all_keys = list(state_dict.keys())
|
||||
for key in all_keys:
|
||||
if "relative_position_index" in key:
|
||||
state_dict.pop(key)
|
||||
|
||||
if "relative_position_bias_table" in key:
|
||||
rel_pos_bias = state_dict[key]
|
||||
src_num_pos, num_attn_heads = rel_pos_bias.size()
|
||||
dst_num_pos, _ = model.visual.state_dict()[key].size()
|
||||
dst_patch_shape = model.visual.patch_embed.patch_shape
|
||||
if dst_patch_shape[0] != dst_patch_shape[1]:
|
||||
raise NotImplementedError()
|
||||
num_extra_tokens = dst_num_pos - (dst_patch_shape[0] * 2 - 1) * (dst_patch_shape[1] * 2 - 1)
|
||||
src_size = int((src_num_pos - num_extra_tokens) ** 0.5)
|
||||
dst_size = int((dst_num_pos - num_extra_tokens) ** 0.5)
|
||||
if src_size != dst_size:
|
||||
print("Position interpolate for %s from %dx%d to %dx%d" % (
|
||||
key, src_size, src_size, dst_size, dst_size))
|
||||
extra_tokens = rel_pos_bias[-num_extra_tokens:, :]
|
||||
rel_pos_bias = rel_pos_bias[:-num_extra_tokens, :]
|
||||
|
||||
def geometric_progression(a, r, n):
|
||||
return a * (1.0 - r ** n) / (1.0 - r)
|
||||
|
||||
left, right = 1.01, 1.5
|
||||
while right - left > 1e-6:
|
||||
q = (left + right) / 2.0
|
||||
gp = geometric_progression(1, q, src_size // 2)
|
||||
if gp > dst_size // 2:
|
||||
right = q
|
||||
else:
|
||||
left = q
|
||||
|
||||
# if q > 1.090307:
|
||||
# q = 1.090307
|
||||
|
||||
dis = []
|
||||
cur = 1
|
||||
for i in range(src_size // 2):
|
||||
dis.append(cur)
|
||||
cur += q ** (i + 1)
|
||||
|
||||
r_ids = [-_ for _ in reversed(dis)]
|
||||
|
||||
x = r_ids + [0] + dis
|
||||
y = r_ids + [0] + dis
|
||||
|
||||
t = dst_size // 2.0
|
||||
dx = np.arange(-t, t + 0.1, 1.0)
|
||||
dy = np.arange(-t, t + 0.1, 1.0)
|
||||
|
||||
print("Original positions = %s" % str(x))
|
||||
print("Target positions = %s" % str(dx))
|
||||
|
||||
all_rel_pos_bias = []
|
||||
|
||||
for i in range(num_attn_heads):
|
||||
z = rel_pos_bias[:, i].view(src_size, src_size).float().numpy()
|
||||
f = F.interpolate.interp2d(x, y, z, kind='cubic')
|
||||
all_rel_pos_bias.append(
|
||||
torch.Tensor(f(dx, dy)).contiguous().view(-1, 1).to(rel_pos_bias.device))
|
||||
|
||||
rel_pos_bias = torch.cat(all_rel_pos_bias, dim=-1)
|
||||
|
||||
new_rel_pos_bias = torch.cat((rel_pos_bias, extra_tokens), dim=0)
|
||||
state_dict[key] = new_rel_pos_bias
|
||||
|
||||
# interpolate position embedding
|
||||
if 'pos_embed' in state_dict:
|
||||
pos_embed_checkpoint = state_dict['pos_embed']
|
||||
embedding_size = pos_embed_checkpoint.shape[-1]
|
||||
num_patches = model.visual.patch_embed.num_patches
|
||||
num_extra_tokens = model.visual.pos_embed.shape[-2] - num_patches
|
||||
# height (== width) for the checkpoint position embedding
|
||||
orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
|
||||
# height (== width) for the new position embedding
|
||||
new_size = int(num_patches ** 0.5)
|
||||
# class_token and dist_token are kept unchanged
|
||||
if orig_size != new_size:
|
||||
print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
|
||||
extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
|
||||
# only the position tokens are interpolated
|
||||
pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
|
||||
pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
|
||||
pos_tokens = torch.nn.functional.interpolate(
|
||||
pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
|
||||
pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
|
||||
new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
|
||||
state_dict['pos_embed'] = new_pos_embed
|
||||
|
||||
patch_embed_proj = state_dict['patch_embed.proj.weight']
|
||||
patch_size = model.visual.patch_embed.patch_size
|
||||
state_dict['patch_embed.proj.weight'] = torch.nn.functional.interpolate(
|
||||
patch_embed_proj.float(), size=patch_size, mode='bicubic', align_corners=False)
|
||||
|
||||
|
||||
def freeze_batch_norm_2d(module, module_match={}, name=''):
|
||||
"""
|
||||
Converts all `BatchNorm2d` and `SyncBatchNorm` layers of provided module into `FrozenBatchNorm2d`. If `module` is
|
||||
itself an instance of either `BatchNorm2d` or `SyncBatchNorm`, it is converted into `FrozenBatchNorm2d` and
|
||||
returned. Otherwise, the module is walked recursively and submodules are converted in place.
|
||||
|
||||
Args:
|
||||
module (torch.nn.Module): Any PyTorch module.
|
||||
module_match (dict): Dictionary of full module names to freeze (all if empty)
|
||||
name (str): Full module name (prefix)
|
||||
|
||||
Returns:
|
||||
torch.nn.Module: Resulting module
|
||||
|
||||
Inspired by https://github.com/pytorch/pytorch/blob/a5895f85be0f10212791145bfedc0261d364f103/torch/nn/modules/batchnorm.py#L762
|
||||
"""
|
||||
res = module
|
||||
is_match = True
|
||||
if module_match:
|
||||
is_match = name in module_match
|
||||
if is_match and isinstance(module, (nn.modules.batchnorm.BatchNorm2d, nn.modules.batchnorm.SyncBatchNorm)):
|
||||
res = FrozenBatchNorm2d(module.num_features)
|
||||
res.num_features = module.num_features
|
||||
res.affine = module.affine
|
||||
if module.affine:
|
||||
res.weight.data = module.weight.data.clone().detach()
|
||||
res.bias.data = module.bias.data.clone().detach()
|
||||
res.running_mean.data = module.running_mean.data
|
||||
res.running_var.data = module.running_var.data
|
||||
res.eps = module.eps
|
||||
else:
|
||||
for child_name, child in module.named_children():
|
||||
full_child_name = '.'.join([name, child_name]) if name else child_name
|
||||
new_child = freeze_batch_norm_2d(child, module_match, full_child_name)
|
||||
if new_child is not child:
|
||||
res.add_module(child_name, new_child)
|
||||
return res
|
||||
|
||||
|
||||
# From PyTorch internals
|
||||
def _ntuple(n):
|
||||
def parse(x):
|
||||
if isinstance(x, collections.abc.Iterable):
|
||||
return x
|
||||
return tuple(repeat(x, n))
|
||||
return parse
|
||||
|
||||
|
||||
to_1tuple = _ntuple(1)
|
||||
to_2tuple = _ntuple(2)
|
||||
to_3tuple = _ntuple(3)
|
||||
to_4tuple = _ntuple(4)
|
||||
to_ntuple = lambda n, x: _ntuple(n)(x)
|
||||
|
||||
|
||||
def is_logging(args):
|
||||
def is_global_master(args):
|
||||
return args.rank == 0
|
||||
|
||||
def is_local_master(args):
|
||||
return args.local_rank == 0
|
||||
|
||||
def is_master(args, local=False):
|
||||
return is_local_master(args) if local else is_global_master(args)
|
||||
return is_master
|
||||
|
||||
|
||||
class AllGather(torch.autograd.Function):
|
||||
"""An autograd function that performs allgather on a tensor.
|
||||
Performs all_gather operation on the provided tensors.
|
||||
*** Warning ***: torch.distributed.all_gather has no gradient.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, tensor, rank, world_size):
|
||||
tensors_gather = [torch.empty_like(tensor) for _ in range(world_size)]
|
||||
torch.distributed.all_gather(tensors_gather, tensor)
|
||||
ctx.rank = rank
|
||||
ctx.batch_size = tensor.shape[0]
|
||||
return torch.cat(tensors_gather, 0)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
return (
|
||||
grad_output[ctx.batch_size * ctx.rank: ctx.batch_size * (ctx.rank + 1)],
|
||||
None,
|
||||
None
|
||||
)
|
||||
|
||||
allgather = AllGather.apply
|
||||
@@ -0,0 +1,378 @@
|
||||
import os
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Tuple
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch import nn, Tensor
|
||||
from transformers import AutoModel, AutoTokenizer, AutoConfig
|
||||
from transformers.file_utils import ModelOutput
|
||||
|
||||
|
||||
from .eva_clip import create_eva_vision_and_transforms
|
||||
from PIL import Image
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EncoderOutput(ModelOutput):
|
||||
q_reps: Optional[Tensor] = None
|
||||
c_reps: Optional[Tensor] = None
|
||||
loss: Optional[Tensor] = None
|
||||
scores: Optional[Tensor] = None
|
||||
|
||||
|
||||
class Visualized_BGE(nn.Module):
|
||||
def __init__(self,
|
||||
model_name_bge: str = None,
|
||||
model_weight = None, # "/path/to/your/weight/file/"
|
||||
normlized: bool = True,
|
||||
sentence_pooling_method: str = 'cls',
|
||||
negatives_cross_device: bool = False,
|
||||
temperature: float = 0.02, # 1.0
|
||||
from_pretrained=None, # local config file and model
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
assert 'bge' in model_name_bge
|
||||
assert model_weight is not None
|
||||
|
||||
self.model_name_bge = model_name_bge
|
||||
|
||||
if 'bge-base-en-v1.5' in model_name_bge:
|
||||
model_name_eva = "EVA02-CLIP-B-16"
|
||||
self.hidden_dim = 768
|
||||
self.depth = 12
|
||||
elif 'bge-m3' in model_name_bge:
|
||||
model_name_eva = "EVA02-CLIP-L-14"
|
||||
self.hidden_dim = 1024
|
||||
self.depth = 24
|
||||
else:
|
||||
raise Exception(f'Unavailable model_name {model_name_bge}')
|
||||
|
||||
if not from_pretrained:
|
||||
bge_config = AutoConfig.from_pretrained(model_name_bge)
|
||||
bge = AutoModel.from_config(bge_config)
|
||||
else:
|
||||
print("Loading from local path.")
|
||||
bge_config = AutoConfig.from_pretrained(from_pretrained, local_files_only=True)
|
||||
bge = AutoModel.from_config(bge_config)
|
||||
|
||||
self.bge_encoder = bge.encoder
|
||||
self.bge_embeddings = bge.embeddings
|
||||
self.bge_pooler = bge.pooler
|
||||
|
||||
self.model_visual, self.preprocess_train, self.preprocess_val= create_eva_vision_and_transforms(
|
||||
model_name_eva,
|
||||
force_custom_clip=True)
|
||||
|
||||
|
||||
self.visual_proj = nn.Linear(self.hidden_dim, self.hidden_dim)
|
||||
|
||||
|
||||
self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')
|
||||
|
||||
self.normlized = normlized
|
||||
self.sentence_pooling_method = sentence_pooling_method
|
||||
self.temperature = temperature
|
||||
if not normlized:
|
||||
self.temperature = 1.0
|
||||
logger.info("reset temperature = 1.0 due to using inner product to compute similarity")
|
||||
|
||||
self.negatives_cross_device = negatives_cross_device
|
||||
if self.negatives_cross_device:
|
||||
if not dist.is_initialized():
|
||||
raise ValueError('Distributed training has not been initialized for representation all gather.')
|
||||
|
||||
self.process_rank = dist.get_rank()
|
||||
self.world_size = dist.get_world_size()
|
||||
|
||||
self.load_model(model_weight)
|
||||
|
||||
if not from_pretrained:
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name_bge, use_fast=False)
|
||||
else:
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(from_pretrained, use_fast=False)
|
||||
|
||||
if torch.cuda.is_available():
|
||||
self.device = torch.device('cuda')
|
||||
self.to(self.device)
|
||||
else:
|
||||
self.device = torch.device('cpu')
|
||||
self.dtype = next(bge.parameters()).dtype
|
||||
|
||||
def load_model(self, model_weight):
|
||||
self.load_state_dict(torch.load(model_weight, map_location='cpu'))
|
||||
|
||||
def gradient_checkpointing_enable(self, **kwargs):
|
||||
# self.bge_encoder.gradient_checkpointing_enable()
|
||||
self.model_visual.set_grad_checkpointing(True)
|
||||
|
||||
|
||||
|
||||
def encode(self, image=None, text=None):
|
||||
# used for simple inference
|
||||
if image is not None:
|
||||
image = self.preprocess_val(Image.open(image)).unsqueeze(0)
|
||||
|
||||
if text is not None:
|
||||
text = self.tokenizer(text, return_tensors="pt", padding=True)
|
||||
return self.encode_mm(image.to(self.device), text.to(self.device))
|
||||
else:
|
||||
return self.encode_image(image.to(self.device))
|
||||
else:
|
||||
if text is not None:
|
||||
text = self.tokenizer(text, return_tensors="pt", padding=True)
|
||||
return self.encode_text(text.to(self.device))
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def get_extended_attention_mask(
|
||||
self, attention_mask: Tensor, input_shape: Tuple[int], device: torch.device = None, dtype: torch.float = torch.float16
|
||||
) -> Tensor:
|
||||
"""
|
||||
Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
|
||||
|
||||
Arguments:
|
||||
attention_mask (`torch.Tensor`):
|
||||
Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
|
||||
input_shape (`Tuple[int]`):
|
||||
The shape of the input to the model.
|
||||
|
||||
Returns:
|
||||
`torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.
|
||||
"""
|
||||
|
||||
# We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
|
||||
# ourselves in which case we just need to make it broadcastable to all heads.
|
||||
if attention_mask.dim() == 3:
|
||||
extended_attention_mask = attention_mask[:, None, :, :]
|
||||
elif attention_mask.dim() == 2:
|
||||
# Provided a padding mask of dimensions [batch_size, seq_length]
|
||||
# - if the model is a decoder, apply a causal mask in addition to the padding mask
|
||||
# - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
|
||||
|
||||
extended_attention_mask = attention_mask[:, None, None, :]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Wrong shape for input_ids (shape {input_shape}) or attention_mask (shape {attention_mask.shape})"
|
||||
)
|
||||
|
||||
# Since attention_mask is 1.0 for positions we want to attend and 0.0 for
|
||||
# masked positions, this operation will create a tensor which is 0.0 for
|
||||
# positions we want to attend and the dtype's smallest value for masked positions.
|
||||
# Since we are adding it to the raw scores before the softmax, this is
|
||||
# effectively the same as removing these entirely.
|
||||
extended_attention_mask = extended_attention_mask.to(dtype=dtype) # fp16 compatibility
|
||||
extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(dtype).min
|
||||
|
||||
return extended_attention_mask
|
||||
|
||||
def sentence_embedding(self, hidden_state, mask):
|
||||
if self.sentence_pooling_method == 'mean':
|
||||
s = torch.sum(hidden_state * mask.unsqueeze(-1).float(), dim=1)
|
||||
d = mask.sum(axis=1, keepdim=True).float()
|
||||
return s / d
|
||||
elif self.sentence_pooling_method == 'cls':
|
||||
return hidden_state[:, 0]
|
||||
|
||||
|
||||
def encode_text(self, texts):
|
||||
'''
|
||||
encode text only
|
||||
'''
|
||||
input_ids = texts['input_ids']
|
||||
attention_mask = texts['attention_mask']
|
||||
|
||||
input_shape = input_ids.size()
|
||||
device = input_ids.device
|
||||
|
||||
token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
|
||||
|
||||
head_mask = [None] * self.depth
|
||||
extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape).to(self.dtype)
|
||||
|
||||
embedding_output = self.bge_embeddings(
|
||||
input_ids=input_ids,
|
||||
position_ids=None,
|
||||
token_type_ids=token_type_ids,
|
||||
inputs_embeds=None,
|
||||
past_key_values_length=0,
|
||||
)
|
||||
encoder_outputs = self.bge_encoder(
|
||||
embedding_output,
|
||||
attention_mask=extended_attention_mask,
|
||||
head_mask=head_mask,
|
||||
encoder_hidden_states=None,
|
||||
encoder_attention_mask=None,
|
||||
past_key_values=None,
|
||||
use_cache=False,
|
||||
output_attentions=False,
|
||||
output_hidden_states=False,
|
||||
return_dict=True,
|
||||
)
|
||||
sequence_output = encoder_outputs[0]
|
||||
# pooled_output = self.bge_pooler(sequence_output) if self.bge_pooler is not None else None
|
||||
|
||||
t_reps = self.sentence_embedding(sequence_output, texts['attention_mask']) # tensor: reps with pooling
|
||||
if self.normlized:
|
||||
t_reps = torch.nn.functional.normalize(t_reps, dim=-1)
|
||||
return t_reps.contiguous()
|
||||
|
||||
def encode_mm(self, images:torch.Tensor, texts):
|
||||
img_token_emb = self.img_token_embedding(images) #[B, Patch_num, C]
|
||||
img_token_emb = img_token_emb[:,1:] # img_cls is not used here
|
||||
img_token_emb = self.visual_proj(img_token_emb)
|
||||
device = img_token_emb.device
|
||||
|
||||
img_token_len = img_token_emb.size()[1]
|
||||
|
||||
# image position embedding, default position: bge_cls + img tokens + texts
|
||||
img_token_position_ids = torch.arange(1, 1 + img_token_len).to(device=device)
|
||||
img_position_embeddings = self.bge_embeddings.position_embeddings(img_token_position_ids)
|
||||
img_token_emb = img_token_emb + img_position_embeddings
|
||||
|
||||
img_token_emb = self.bge_embeddings.LayerNorm(img_token_emb)
|
||||
|
||||
### deal with prompt/text
|
||||
prompt_input_ids = texts['input_ids']
|
||||
prompt_attention_mask = texts['attention_mask']
|
||||
prom_input_shape = prompt_input_ids.size()
|
||||
|
||||
# bert
|
||||
batch_size = prom_input_shape[0]
|
||||
prompt_len = prom_input_shape[1]
|
||||
prompt_start = 1 + img_token_len
|
||||
|
||||
|
||||
cls_id = torch.tensor([0]).to(device=device)
|
||||
prompt_position_ids = torch.arange(prompt_start, prompt_start + prompt_len - 1).to(device=device)
|
||||
prompt_position_ids = torch.cat([cls_id, prompt_position_ids]).to(device=device)
|
||||
|
||||
prompt_token_type_ids = torch.zeros(prom_input_shape, dtype=torch.long, device=device)
|
||||
prompt_embedding_output = self.bge_embeddings(
|
||||
input_ids=prompt_input_ids,
|
||||
position_ids=prompt_position_ids,
|
||||
token_type_ids=prompt_token_type_ids,
|
||||
inputs_embeds=None,
|
||||
past_key_values_length=0,
|
||||
) # [B, T, C]
|
||||
|
||||
|
||||
cls_token = prompt_embedding_output[:, 0:1, :] # bge_cls token
|
||||
prompt_embedding_output = prompt_embedding_output[:, 1:]
|
||||
|
||||
prompt_img_embedding = torch.cat([cls_token, img_token_emb, prompt_embedding_output], dim=1)
|
||||
|
||||
img_attention_mask = torch.ones(batch_size, img_token_len, device=device)
|
||||
prom_img_attention_mask = torch.cat([img_attention_mask, prompt_attention_mask], dim=1)
|
||||
prom_img_input_shape = prompt_img_embedding.size()
|
||||
|
||||
head_mask = [None] * self.depth
|
||||
extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(prom_img_attention_mask, prom_img_input_shape).to(self.dtype)
|
||||
|
||||
|
||||
encoder_outputs = self.bge_encoder(
|
||||
prompt_img_embedding,
|
||||
attention_mask=extended_attention_mask,
|
||||
head_mask=head_mask,
|
||||
encoder_hidden_states=None,
|
||||
encoder_attention_mask=None,
|
||||
past_key_values=None,
|
||||
use_cache=False,
|
||||
output_attentions=False,
|
||||
output_hidden_states=False,
|
||||
return_dict=True,
|
||||
)
|
||||
sequence_output = encoder_outputs[0]
|
||||
|
||||
prompt_img_reps = self.sentence_embedding(sequence_output, prom_img_attention_mask) # tensor: reps with pooling
|
||||
if self.normlized:
|
||||
prompt_img_reps = torch.nn.functional.normalize(prompt_img_reps, dim=-1)
|
||||
return prompt_img_reps
|
||||
|
||||
def compute_similarity(self, q_reps, p_reps):
|
||||
if len(p_reps.size()) == 2:
|
||||
return torch.matmul(q_reps, p_reps.transpose(0, 1))
|
||||
return torch.matmul(q_reps, p_reps.transpose(-2, -1))
|
||||
|
||||
def img_token_embedding(self, images):
|
||||
if images is None:
|
||||
return None
|
||||
img_token_emb = self.model_visual.encode_image(images, normalize=False) # return_all_features=True, [B, Patch_num, C]
|
||||
|
||||
return img_token_emb.contiguous()
|
||||
|
||||
def encode_image(self, images):
|
||||
if images is None:
|
||||
return None
|
||||
|
||||
batch_size = images.shape[0]
|
||||
prompts = [""] * batch_size
|
||||
|
||||
prompts = self.tokenizer(prompts, return_tensors="pt", padding=True)
|
||||
prompts = prompts.to(images.device)
|
||||
img_reps = self.encode_mm(images, prompts)
|
||||
return img_reps
|
||||
|
||||
def forward(self, mm_it_query=None, image_candidate=None, text_candidate=None, text_query=None, mm_it_candidate=None, task_type=None):
|
||||
### for stage-2 training
|
||||
if task_type == "edit_image":
|
||||
mm_query_reps = self.encode_mm(mm_it_query[0], mm_it_query[1])
|
||||
image_candi_reps = self.encode_image(image_candidate)
|
||||
query_reps = mm_query_reps
|
||||
candi_reps = image_candi_reps
|
||||
|
||||
elif task_type == "t2it":
|
||||
text_query_reps = self.encode_text(text_query)
|
||||
mmit_candi_reps = self.encode_mm(mm_it_candidate[0], mm_it_candidate[1])
|
||||
query_reps = text_query_reps
|
||||
candi_reps = mmit_candi_reps
|
||||
|
||||
|
||||
if self.training:
|
||||
if self.negatives_cross_device:
|
||||
query_reps = self._dist_gather_tensor(query_reps)
|
||||
candi_reps = self._dist_gather_tensor(candi_reps)
|
||||
|
||||
scores = self.compute_similarity(query_reps, candi_reps)
|
||||
scores = scores / self.temperature
|
||||
scores = scores.view(query_reps.size(0), -1)
|
||||
|
||||
target = torch.arange(scores.size(0), device=scores.device, dtype=torch.long)
|
||||
target = target * (candi_reps.size(0) // query_reps.size(0))
|
||||
|
||||
loss_edit = self.compute_loss(scores, target)
|
||||
loss = loss_edit
|
||||
|
||||
logging.info("task types: %s; loss: %s" %(task_type, str(loss_edit)))
|
||||
else:
|
||||
scores = self.compute_similarity(query_reps, candi_reps)
|
||||
loss=None
|
||||
return EncoderOutput(
|
||||
loss=loss,
|
||||
scores=scores,
|
||||
q_reps=query_reps,
|
||||
c_reps=candi_reps,
|
||||
)
|
||||
|
||||
def compute_loss(self, scores, target):
|
||||
return self.cross_entropy(scores, target)
|
||||
|
||||
def _dist_gather_tensor(self, t: Optional[torch.Tensor]):
|
||||
if t is None:
|
||||
return None
|
||||
t = t.contiguous()
|
||||
|
||||
all_tensors = [torch.empty_like(t) for _ in range(self.world_size)]
|
||||
dist.all_gather(all_tensors, t)
|
||||
|
||||
all_tensors[self.process_rank] = t
|
||||
all_tensors = torch.cat(all_tensors, dim=0)
|
||||
|
||||
return all_tensors
|
||||
|
||||
def save(self, output_dir: str):
|
||||
torch.save(self.state_dict(), os.path.join(output_dir, 'Visualized_BGE.pth'))
|
||||
@@ -0,0 +1,403 @@
|
||||
import json
|
||||
import os
|
||||
from tqdm import tqdm
|
||||
import torch
|
||||
from visual_bge.visual_bge.modeling import Visualized_BGE
|
||||
from pymilvus import connections, MilvusClient, FieldSchema, CollectionSchema, DataType, Collection, AnnSearchRequest, RRFRanker
|
||||
from pymilvus.model.hybrid import BGEM3EmbeddingFunction
|
||||
import numpy as np
|
||||
import cv2
|
||||
from PIL import Image
|
||||
from typing import List, Dict, Any
|
||||
from dataclasses import dataclass
|
||||
|
||||
@dataclass
|
||||
class DragonImage:
|
||||
"""龙类图像数据类"""
|
||||
img_id: str
|
||||
path: str
|
||||
title: str
|
||||
description: str
|
||||
category: str
|
||||
location: str
|
||||
environment: str
|
||||
combat_details: Dict[str, Any] = None
|
||||
scene_info: Dict[str, Any] = None
|
||||
|
||||
class DragonDataset:
|
||||
"""龙类图像数据集管理类"""
|
||||
def __init__(self, data_dir: str, metadata_path: str):
|
||||
self.data_dir = data_dir
|
||||
self.metadata_path = metadata_path
|
||||
self.images: List[DragonImage] = []
|
||||
self._load_metadata()
|
||||
|
||||
def _load_metadata(self):
|
||||
"""加载图像元数据"""
|
||||
with open(self.metadata_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
for img_data in data:
|
||||
# 确保图片路径是完整的
|
||||
if not img_data['path'].startswith(self.data_dir):
|
||||
img_data['path'] = os.path.join(self.data_dir, img_data['path'].split('/')[-1])
|
||||
self.images.append(DragonImage(**img_data))
|
||||
|
||||
def get_text_content(self, img: DragonImage) -> str:
|
||||
"""获取图像的文本描述内容"""
|
||||
parts = [
|
||||
img.title, img.description,
|
||||
img.location, img.environment
|
||||
]
|
||||
if img.combat_details:
|
||||
parts.extend(img.combat_details.get('combat_style', []))
|
||||
parts.extend(img.combat_details.get('abilities_used', []))
|
||||
if img.scene_info:
|
||||
parts.append(img.scene_info.get('time_of_day', ''))
|
||||
return ' '.join(filter(None, parts))
|
||||
|
||||
class HybridMultimodalEncoder:
|
||||
"""混合多模态编码器类"""
|
||||
def __init__(self, visual_model_name: str, visual_model_path: str):
|
||||
# 初始化Visual-BGE模型（用于多模态）
|
||||
self.visual_model = Visualized_BGE(model_name_bge=visual_model_name, model_weight=visual_model_path)
|
||||
self.visual_model.eval()
|
||||
|
||||
# 初始化BGE-M3模型（用于混合检索）
|
||||
self.bge_m3 = BGEM3EmbeddingFunction(use_fp16=False, device="cpu")
|
||||
print(f"BGE-M3 密集向量维度: {self.bge_m3.dim['dense']}")
|
||||
|
||||
def encode_multimodal(self, image_path: str, text: str) -> list[float]:
|
||||
"""编码多模态内容（图像+文本）"""
|
||||
with torch.no_grad():
|
||||
query_emb = self.visual_model.encode(image=image_path, text=text)
|
||||
return query_emb.tolist()[0]
|
||||
|
||||
def encode_text_hybrid(self, text: str) -> dict:
|
||||
"""使用BGE-M3编码文本，返回稀疏和密集向量"""
|
||||
embeddings = self.bge_m3([text])
|
||||
return {
|
||||
'sparse': embeddings["sparse"],
|
||||
'dense': embeddings["dense"]
|
||||
}
|
||||
|
||||
def encode_query(self, image_path: str = None, text: str = None, mode: str = "multimodal") -> dict:
|
||||
"""编码查询，支持多种模式"""
|
||||
result = {}
|
||||
|
||||
if mode in ["multimodal", "all"] and image_path and text:
|
||||
result['multimodal'] = self.encode_multimodal(image_path, text)
|
||||
|
||||
if mode in ["hybrid", "dense", "sparse", "all"] and text:
|
||||
text_embeddings = self.encode_text_hybrid(text)
|
||||
result['dense'] = text_embeddings['dense'][0]
|
||||
result['sparse'] = text_embeddings['sparse']._getrow(0)
|
||||
|
||||
return result
|
||||
|
||||
def visualize_results(query_image_path: str, retrieved_results: list, search_mode: str,
|
||||
img_height: int = 300, img_width: int = 300, row_count: int = 3) -> np.ndarray:
|
||||
"""从检索到的结果创建一个全景图用于可视化"""
|
||||
panoramic_width = img_width * row_count
|
||||
panoramic_height = img_height * row_count
|
||||
panoramic_image = np.full((panoramic_height, panoramic_width, 3), 255, dtype=np.uint8)
|
||||
query_display_area = np.full((panoramic_height, img_width, 3), 255, dtype=np.uint8)
|
||||
|
||||
# 处理查询图像
|
||||
if query_image_path and os.path.exists(query_image_path):
|
||||
query_pil = Image.open(query_image_path).convert("RGB")
|
||||
query_cv = np.array(query_pil)[:, :, ::-1]
|
||||
resized_query = cv2.resize(query_cv, (img_width, img_height))
|
||||
bordered_query = cv2.copyMakeBorder(resized_query, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=(255, 0, 0))
|
||||
query_display_area[img_height * (row_count - 1):, :] = cv2.resize(bordered_query, (img_width, img_height))
|
||||
cv2.putText(query_display_area, "Query", (10, panoramic_height - 40), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 0, 0), 2)
|
||||
cv2.putText(query_display_area, search_mode, (10, panoramic_height - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 100, 0), 2)
|
||||
|
||||
# 处理检索到的图像
|
||||
for i, result in enumerate(retrieved_results):
|
||||
row, col = i // row_count, i % row_count
|
||||
start_row, start_col = row * img_height, col * img_width
|
||||
|
||||
img_path = result['image_path']
|
||||
retrieved_pil = Image.open(img_path).convert("RGB")
|
||||
retrieved_cv = np.array(retrieved_pil)[:, :, ::-1]
|
||||
resized_retrieved = cv2.resize(retrieved_cv, (img_width - 4, img_height - 4))
|
||||
bordered_retrieved = cv2.copyMakeBorder(resized_retrieved, 2, 2, 2, 2, cv2.BORDER_CONSTANT, value=(0, 0, 0))
|
||||
panoramic_image[start_row:start_row + img_height, start_col:start_col + img_width] = bordered_retrieved
|
||||
|
||||
# 添加索引号和相似度
|
||||
cv2.putText(panoramic_image, f"{i+1}", (start_col + 10, start_row + 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
|
||||
cv2.putText(panoramic_image, f"{result['distance']:.3f}", (start_col + 10, start_row + img_height - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
|
||||
|
||||
return np.hstack([query_display_area, panoramic_image])
|
||||
|
||||
class HybridMultimodalSearcher:
|
||||
"""混合多模态搜索系统"""
|
||||
def __init__(self, data_dir: str, metadata_path: str, collection_name: str, milvus_uri: str):
|
||||
self.data_dir = data_dir
|
||||
self.metadata_path = metadata_path
|
||||
self.collection_name = collection_name
|
||||
self.milvus_uri = milvus_uri
|
||||
|
||||
# 初始化数据集和编码器
|
||||
print("--> 正在初始化数据集...")
|
||||
self.dataset = DragonDataset(data_dir, metadata_path)
|
||||
print(f"加载了 {len(self.dataset.images)} 张龙类图像")
|
||||
|
||||
print("--> 正在初始化混合多模态编码器...")
|
||||
self.encoder = HybridMultimodalEncoder(
|
||||
visual_model_name="BAAI/bge-base-en-v1.5",
|
||||
visual_model_path="../../models/bge/Visualized_base_en_v1.5.pth"
|
||||
)
|
||||
|
||||
# 连接Milvus
|
||||
print(f"--> 正在连接到 Milvus: {milvus_uri}")
|
||||
connections.connect(uri=milvus_uri)
|
||||
self.milvus_client = MilvusClient(uri=milvus_uri)
|
||||
|
||||
self.collection = None
|
||||
|
||||
def create_collection(self):
|
||||
"""创建Collection"""
|
||||
print(f"--> 正在创建 Collection '{self.collection_name}'")
|
||||
if self.milvus_client.has_collection(self.collection_name):
|
||||
self.milvus_client.drop_collection(self.collection_name)
|
||||
print(f"已删除已存在的 Collection: '{self.collection_name}'")
|
||||
|
||||
# 获取向量维度
|
||||
sample_text = self.dataset.get_text_content(self.dataset.images[0])
|
||||
sample_path = self.dataset.images[0].path
|
||||
multimodal_dim = len(self.encoder.encode_multimodal(sample_path, sample_text))
|
||||
dense_dim = self.encoder.bge_m3.dim["dense"]
|
||||
|
||||
print(f"多模态向量维度: {multimodal_dim}")
|
||||
print(f"密集向量维度: {dense_dim}")
|
||||
|
||||
fields = [
|
||||
FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100),
|
||||
FieldSchema(name="img_id", dtype=DataType.VARCHAR, max_length=100),
|
||||
FieldSchema(name="image_path", dtype=DataType.VARCHAR, max_length=512),
|
||||
FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
|
||||
FieldSchema(name="description", dtype=DataType.VARCHAR, max_length=4096),
|
||||
FieldSchema(name="category", dtype=DataType.VARCHAR, max_length=64),
|
||||
FieldSchema(name="location", dtype=DataType.VARCHAR, max_length=128),
|
||||
FieldSchema(name="environment", dtype=DataType.VARCHAR, max_length=64),
|
||||
# 三种向量类型
|
||||
FieldSchema(name="multimodal_vector", dtype=DataType.FLOAT_VECTOR, dim=multimodal_dim),
|
||||
FieldSchema(name="text_sparse_vector", dtype=DataType.SPARSE_FLOAT_VECTOR),
|
||||
FieldSchema(name="text_dense_vector", dtype=DataType.FLOAT_VECTOR, dim=dense_dim)
|
||||
]
|
||||
|
||||
schema = CollectionSchema(fields, description="混合多模态龙类图像检索")
|
||||
self.collection = Collection(name=self.collection_name, schema=schema, consistency_level="Strong")
|
||||
print("--> Collection 创建成功")
|
||||
|
||||
# 创建索引
|
||||
print("--> 正在创建索引...")
|
||||
# 多模态向量索引
|
||||
multimodal_index = {"index_type": "HNSW", "metric_type": "COSINE", "params": {"M": 16, "efConstruction": 256}}
|
||||
self.collection.create_index("multimodal_vector", multimodal_index)
|
||||
print("多模态向量索引创建成功")
|
||||
|
||||
# 稀疏向量索引
|
||||
sparse_index = {"index_type": "SPARSE_INVERTED_INDEX", "metric_type": "IP"}
|
||||
self.collection.create_index("text_sparse_vector", sparse_index)
|
||||
print("稀疏向量索引创建成功")
|
||||
|
||||
# 密集向量索引
|
||||
dense_index = {"index_type": "AUTOINDEX", "metric_type": "IP"}
|
||||
self.collection.create_index("text_dense_vector", dense_index)
|
||||
print("密集向量索引创建成功")
|
||||
|
||||
self.collection.load()
|
||||
print(f"--> Collection '{self.collection_name}' 已加载到内存")
|
||||
|
||||
def insert_data(self):
|
||||
"""插入数据"""
|
||||
if self.collection.is_empty:
|
||||
print("--> Collection 为空，开始插入数据...")
|
||||
|
||||
# 准备批量数据
|
||||
img_ids, image_paths, titles, descriptions = [], [], [], []
|
||||
categories, locations, environments = [], [], []
|
||||
multimodal_vectors, text_sparse_vectors, text_dense_vectors = [], [], []
|
||||
|
||||
for img_data in tqdm(self.dataset.images, desc="生成向量嵌入"):
|
||||
text_content = self.dataset.get_text_content(img_data)
|
||||
|
||||
# 生成多模态向量（图像+文本）
|
||||
multimodal_vector = self.encoder.encode_multimodal(img_data.path, text_content)
|
||||
|
||||
# 生成文本的混合向量（稀疏+密集）
|
||||
text_embeddings = self.encoder.encode_text_hybrid(text_content)
|
||||
|
||||
# 收集数据
|
||||
img_ids.append(img_data.img_id)
|
||||
image_paths.append(img_data.path)
|
||||
titles.append(img_data.title)
|
||||
descriptions.append(img_data.description)
|
||||
categories.append(img_data.category)
|
||||
locations.append(img_data.location)
|
||||
environments.append(img_data.environment)
|
||||
|
||||
multimodal_vectors.append(multimodal_vector)
|
||||
text_sparse_vectors.append(text_embeddings['sparse']._getrow(0))
|
||||
text_dense_vectors.append(text_embeddings['dense'][0])
|
||||
|
||||
# 插入数据
|
||||
self.collection.insert([
|
||||
img_ids, image_paths, titles, descriptions, categories, locations, environments,
|
||||
multimodal_vectors, text_sparse_vectors, text_dense_vectors
|
||||
])
|
||||
|
||||
self.collection.flush()
|
||||
print(f"--> 数据插入完成，总数: {self.collection.num_entities}")
|
||||
else:
|
||||
print(f"--> Collection 中已有 {self.collection.num_entities} 条数据，跳过插入")
|
||||
|
||||
def search(self, query_image_path: str, query_text: str, mode: str = "hybrid", top_k: int = 5) -> list:
|
||||
"""执行搜索"""
|
||||
search_params = {"metric_type": "IP", "params": {}}
|
||||
cosine_params = {"metric_type": "COSINE", "params": {"ef": 128}}
|
||||
output_fields = ["img_id", "image_path", "title", "description", "category", "location", "environment"]
|
||||
|
||||
if mode == "multimodal":
|
||||
# 多模态检索
|
||||
query_vector = self.encoder.encode_multimodal(query_image_path, query_text)
|
||||
results = self.collection.search(
|
||||
[query_vector], "multimodal_vector", param=cosine_params,
|
||||
limit=top_k, output_fields=output_fields
|
||||
)[0]
|
||||
|
||||
elif mode == "dense":
|
||||
# 密集向量检索
|
||||
query_embeddings = self.encoder.encode_text_hybrid(query_text)
|
||||
dense_vec = query_embeddings['dense'][0]
|
||||
results = self.collection.search(
|
||||
[dense_vec], "text_dense_vector", param=search_params,
|
||||
limit=top_k, output_fields=output_fields
|
||||
)[0]
|
||||
|
||||
elif mode == "sparse":
|
||||
# 稀疏向量检索
|
||||
query_embeddings = self.encoder.encode_text_hybrid(query_text)
|
||||
sparse_vec = query_embeddings['sparse']._getrow(0)
|
||||
results = self.collection.search(
|
||||
[sparse_vec], "text_sparse_vector", param=search_params,
|
||||
limit=top_k, output_fields=output_fields
|
||||
)[0]
|
||||
|
||||
elif mode == "hybrid":
|
||||
# 混合检索（稀疏+密集）
|
||||
query_embeddings = self.encoder.encode_text_hybrid(query_text)
|
||||
dense_vec = query_embeddings['dense'][0]
|
||||
sparse_vec = query_embeddings['sparse']._getrow(0)
|
||||
|
||||
# 创建RRF融合器
|
||||
rerank = RRFRanker(k=60)
|
||||
|
||||
# 创建搜索请求
|
||||
dense_req = AnnSearchRequest([dense_vec], "text_dense_vector", search_params, limit=top_k)
|
||||
sparse_req = AnnSearchRequest([sparse_vec], "text_sparse_vector", search_params, limit=top_k)
|
||||
|
||||
# 执行混合搜索
|
||||
results = self.collection.hybrid_search(
|
||||
[sparse_req, dense_req], rerank=rerank, limit=top_k, output_fields=output_fields
|
||||
)[0]
|
||||
|
||||
return results
|
||||
|
||||
def compare_search_modes(self, query_image_path: str, query_text: str, top_k: int = 5):
|
||||
"""对比不同搜索模式的效果"""
|
||||
modes = ["multimodal", "dense", "sparse", "hybrid"]
|
||||
results = {}
|
||||
|
||||
print(f"\n{'='*50}")
|
||||
print(f"查询图像: {query_image_path}")
|
||||
print(f"查询文本: {query_text}")
|
||||
print(f"{'='*50}")
|
||||
|
||||
for mode in modes:
|
||||
print(f"\n--- [{mode.upper()}] 搜索结果 ---")
|
||||
search_results = self.search(query_image_path, query_text, mode, top_k)
|
||||
|
||||
mode_results = []
|
||||
for i, hit in enumerate(search_results):
|
||||
print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
|
||||
print(f" 路径: {hit.entity.get('image_path')}")
|
||||
print(f" 描述: {hit.entity.get('description')[:80]}...")
|
||||
|
||||
mode_results.append({
|
||||
'image_path': hit.entity.get('image_path'),
|
||||
'distance': hit.distance,
|
||||
'title': hit.entity.get('title')
|
||||
})
|
||||
|
||||
results[mode] = mode_results
|
||||
|
||||
return results
|
||||
|
||||
def visualize_comparison(self, query_image_path: str, query_text: str, top_k: int = 3):
|
||||
"""可视化对比不同搜索模式"""
|
||||
modes = ["multimodal", "dense", "sparse", "hybrid"]
|
||||
|
||||
for mode in modes:
|
||||
results = self.search(query_image_path, query_text, mode, top_k)
|
||||
|
||||
retrieved_results = []
|
||||
for hit in results:
|
||||
retrieved_results.append({
|
||||
'image_path': hit.entity.get('image_path'),
|
||||
'distance': hit.distance
|
||||
})
|
||||
|
||||
if retrieved_results:
|
||||
panoramic_image = visualize_results(query_image_path, retrieved_results, mode.upper())
|
||||
output_path = f"../../data/C4/{mode}_search_result.png"
|
||||
cv2.imwrite(output_path, panoramic_image)
|
||||
print(f"{mode.upper()} 搜索结果已保存到: {output_path}")
|
||||
|
||||
def cleanup(self):
|
||||
"""清理资源"""
|
||||
if self.collection:
|
||||
self.collection.release()
|
||||
print(f"已从内存中释放 Collection: '{self.collection_name}'")
|
||||
self.milvus_client.drop_collection(self.collection_name)
|
||||
print(f"已删除 Collection: '{self.collection_name}'")
|
||||
|
||||
# 主程序
|
||||
if __name__ == "__main__":
|
||||
# 初始化设置
|
||||
DATA_DIR = "../../data/C3/dragon"
|
||||
METADATA_PATH = "../../data/C4/metadata/dragon.json"
|
||||
COLLECTION_NAME = "hybrid_multimodal_dragon_demo"
|
||||
MILVUS_URI = "http://localhost:19530"
|
||||
|
||||
# 创建混合多模态搜索系统
|
||||
searcher = HybridMultimodalSearcher(DATA_DIR, METADATA_PATH, COLLECTION_NAME, MILVUS_URI)
|
||||
|
||||
try:
|
||||
# 创建Collection并插入数据
|
||||
searcher.create_collection()
|
||||
searcher.insert_data()
|
||||
|
||||
# 执行搜索对比
|
||||
query_image_path = os.path.join(DATA_DIR, "query.png")
|
||||
query_text = "悬崖上的巨龙"
|
||||
|
||||
# 对比不同搜索模式
|
||||
results = searcher.compare_search_modes(query_image_path, query_text, top_k=3)
|
||||
|
||||
# 可视化结果
|
||||
searcher.visualize_comparison(query_image_path, query_text, top_k=3)
|
||||
|
||||
print(f"\n{'='*50}")
|
||||
print("搜索模式分析:")
|
||||
print("- MULTIMODAL: 结合图像和文本的多模态向量检索")
|
||||
print("- DENSE: 基于语义的密集向量检索")
|
||||
print("- SPARSE: 基于关键词的稀疏向量检索")
|
||||
print("- HYBRID: 稀疏+密集向量RRF融合检索")
|
||||
print(f"{'='*50}")
|
||||
|
||||
finally:
|
||||
# 清理资源
|
||||
searcher.cleanup()
|
||||
@@ -0,0 +1,301 @@
|
||||
import json
|
||||
import os
|
||||
from tqdm import tqdm
|
||||
from glob import glob
|
||||
import torch
|
||||
from visual_bge.visual_bge.modeling import Visualized_BGE
|
||||
from pymilvus import MilvusClient, FieldSchema, CollectionSchema, DataType
|
||||
import numpy as np
|
||||
import cv2
|
||||
from PIL import Image
|
||||
from typing import List, Dict, Any
|
||||
from dataclasses import dataclass
|
||||
|
||||
@dataclass
|
||||
class DragonImage:
|
||||
"""龙类图像数据类"""
|
||||
img_id: str
|
||||
path: str
|
||||
title: str
|
||||
description: str
|
||||
category: str
|
||||
location: str
|
||||
environment: str
|
||||
combat_details: Dict[str, Any] = None
|
||||
scene_info: Dict[str, Any] = None
|
||||
|
||||
class DragonDataset:
|
||||
"""龙类图像数据集管理类"""
|
||||
def __init__(self, data_dir: str, metadata_path: str):
|
||||
self.data_dir = data_dir
|
||||
self.metadata_path = metadata_path
|
||||
self.images: List[DragonImage] = []
|
||||
self._load_metadata()
|
||||
|
||||
def _load_metadata(self):
|
||||
"""加载图像元数据"""
|
||||
with open(self.metadata_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
for img_data in data:
|
||||
# 确保图片路径是完整的
|
||||
if not img_data['path'].startswith(self.data_dir):
|
||||
img_data['path'] = os.path.join(self.data_dir, img_data['path'].split('/')[-1])
|
||||
self.images.append(DragonImage(**img_data))
|
||||
|
||||
def get_image_paths(self) -> List[str]:
|
||||
"""获取所有图像路径"""
|
||||
return [img.path for img in self.images]
|
||||
|
||||
def get_metadata_by_path(self, path: str) -> DragonImage:
|
||||
"""根据路径获取元数据"""
|
||||
for img in self.images:
|
||||
if img.path == path:
|
||||
return img
|
||||
return None
|
||||
|
||||
def get_text_content(self, img: DragonImage) -> str:
|
||||
"""获取图像的文本描述内容"""
|
||||
parts = [
|
||||
img.title, img.description,
|
||||
img.location, img.environment
|
||||
]
|
||||
if img.combat_details:
|
||||
parts.extend(img.combat_details.get('combat_style', []))
|
||||
parts.extend(img.combat_details.get('abilities_used', []))
|
||||
if img.scene_info:
|
||||
parts.append(img.scene_info.get('time_of_day', ''))
|
||||
return ' '.join(filter(None, parts))
|
||||
|
||||
class Encoder:
|
||||
"""编码器类，用于将图像和文本编码为向量"""
|
||||
def __init__(self, model_name: str, model_path: str):
|
||||
self.model = Visualized_BGE(model_name_bge=model_name, model_weight=model_path)
|
||||
self.model.eval()
|
||||
|
||||
def encode_query(self, image_path: str = None, text: str = None) -> list[float]:
|
||||
"""编码查询（支持图像+文本或仅文本）"""
|
||||
with torch.no_grad():
|
||||
if image_path and text:
|
||||
query_emb = self.model.encode(image=image_path, text=text)
|
||||
elif image_path:
|
||||
query_emb = self.model.encode(image=image_path)
|
||||
elif text:
|
||||
query_emb = self.model.encode(text=text)
|
||||
else:
|
||||
raise ValueError("必须提供图像路径或文本内容")
|
||||
return query_emb.tolist()[0]
|
||||
|
||||
def encode_multimodal(self, image_path: str, text: str) -> list[float]:
|
||||
"""编码多模态内容（图像+文本）"""
|
||||
with torch.no_grad():
|
||||
query_emb = self.model.encode(image=image_path, text=text)
|
||||
return query_emb.tolist()[0]
|
||||
|
||||
def visualize_results(query_image_path: str, retrieved_results: list, img_height: int = 300, img_width: int = 300, row_count: int = 3) -> np.ndarray:
|
||||
"""从检索到的结果创建一个全景图用于可视化"""
|
||||
panoramic_width = img_width * row_count
|
||||
panoramic_height = img_height * row_count
|
||||
panoramic_image = np.full((panoramic_height, panoramic_width, 3), 255, dtype=np.uint8)
|
||||
query_display_area = np.full((panoramic_height, img_width, 3), 255, dtype=np.uint8)
|
||||
|
||||
# 处理查询图像
|
||||
if query_image_path and os.path.exists(query_image_path):
|
||||
query_pil = Image.open(query_image_path).convert("RGB")
|
||||
query_cv = np.array(query_pil)[:, :, ::-1]
|
||||
resized_query = cv2.resize(query_cv, (img_width, img_height))
|
||||
bordered_query = cv2.copyMakeBorder(resized_query, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=(255, 0, 0))
|
||||
query_display_area[img_height * (row_count - 1):, :] = cv2.resize(bordered_query, (img_width, img_height))
|
||||
cv2.putText(query_display_area, "Query", (10, panoramic_height - 20), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
|
||||
|
||||
# 处理检索到的图像
|
||||
for i, result in enumerate(retrieved_results):
|
||||
row, col = i // row_count, i % row_count
|
||||
start_row, start_col = row * img_height, col * img_width
|
||||
|
||||
img_path = result['image_path']
|
||||
retrieved_pil = Image.open(img_path).convert("RGB")
|
||||
retrieved_cv = np.array(retrieved_pil)[:, :, ::-1]
|
||||
resized_retrieved = cv2.resize(retrieved_cv, (img_width - 4, img_height - 4))
|
||||
bordered_retrieved = cv2.copyMakeBorder(resized_retrieved, 2, 2, 2, 2, cv2.BORDER_CONSTANT, value=(0, 0, 0))
|
||||
panoramic_image[start_row:start_row + img_height, start_col:start_col + img_width] = bordered_retrieved
|
||||
|
||||
# 添加索引号和相似度
|
||||
cv2.putText(panoramic_image, f"{i+1}", (start_col + 10, start_row + 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
|
||||
cv2.putText(panoramic_image, f"{result['distance']:.3f}", (start_col + 10, start_row + img_height - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
|
||||
|
||||
return np.hstack([query_display_area, panoramic_image])
|
||||
|
||||
# 1. 初始化设置
|
||||
MODEL_NAME = "BAAI/bge-base-en-v1.5"
|
||||
MODEL_PATH = "../../models/bge/Visualized_base_en_v1.5.pth"
|
||||
DATA_DIR = "../../data/C3/dragon"
|
||||
METADATA_PATH = "../../data/C4/metadata/dragon.json"
|
||||
COLLECTION_NAME = "multimodal_dragon_demo"
|
||||
MILVUS_URI = "http://localhost:19530"
|
||||
|
||||
# 2. 初始化数据集和编码器
|
||||
print("--> 正在初始化数据集...")
|
||||
dataset = DragonDataset(DATA_DIR, METADATA_PATH)
|
||||
print(f"加载了 {len(dataset.images)} 张龙类图像")
|
||||
|
||||
print("--> 正在初始化编码器和Milvus客户端...")
|
||||
encoder = Encoder(MODEL_NAME, MODEL_PATH)
|
||||
milvus_client = MilvusClient(uri=MILVUS_URI)
|
||||
|
||||
# 3. 创建 Milvus Collection
|
||||
print(f"\n--> 正在创建 Collection '{COLLECTION_NAME}'")
|
||||
if milvus_client.has_collection(COLLECTION_NAME):
|
||||
milvus_client.drop_collection(COLLECTION_NAME)
|
||||
print(f"已删除已存在的 Collection: '{COLLECTION_NAME}'")
|
||||
|
||||
# 获取向量维度
|
||||
sample_text = dataset.get_text_content(dataset.images[0])
|
||||
sample_path = dataset.images[0].path
|
||||
dim = len(encoder.encode_multimodal(sample_path, sample_text))
|
||||
|
||||
fields = [
|
||||
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
|
||||
FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
||||
FieldSchema(name="img_id", dtype=DataType.VARCHAR, max_length=100),
|
||||
FieldSchema(name="image_path", dtype=DataType.VARCHAR, max_length=512),
|
||||
FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
|
||||
FieldSchema(name="description", dtype=DataType.VARCHAR, max_length=4096),
|
||||
FieldSchema(name="category", dtype=DataType.VARCHAR, max_length=64),
|
||||
FieldSchema(name="location", dtype=DataType.VARCHAR, max_length=128),
|
||||
FieldSchema(name="environment", dtype=DataType.VARCHAR, max_length=64),
|
||||
]
|
||||
|
||||
schema = CollectionSchema(fields, description="多模态龙类图像检索")
|
||||
|
||||
# 创建集合
|
||||
milvus_client.create_collection(collection_name=COLLECTION_NAME, schema=schema)
|
||||
print(f"成功创建 Collection: '{COLLECTION_NAME}'")
|
||||
|
||||
# 4. 准备并插入数据
|
||||
print(f"\n--> 正在向 '{COLLECTION_NAME}' 插入数据")
|
||||
data_to_insert = []
|
||||
for img_data in tqdm(dataset.images, desc="生成多模态嵌入"):
|
||||
# 结合图像和文本信息生成向量
|
||||
text_content = dataset.get_text_content(img_data)
|
||||
vector = encoder.encode_multimodal(img_data.path, text_content)
|
||||
|
||||
data_to_insert.append({
|
||||
"vector": vector,
|
||||
"img_id": img_data.img_id,
|
||||
"image_path": img_data.path,
|
||||
"title": img_data.title,
|
||||
"description": img_data.description,
|
||||
"category": img_data.category,
|
||||
"location": img_data.location,
|
||||
"environment": img_data.environment
|
||||
})
|
||||
|
||||
if data_to_insert:
|
||||
result = milvus_client.insert(collection_name=COLLECTION_NAME, data=data_to_insert)
|
||||
print(f"成功插入 {result['insert_count']} 条数据。")
|
||||
|
||||
# 5. 创建索引
|
||||
print(f"\n--> 正在为 '{COLLECTION_NAME}' 创建索引")
|
||||
index_params = milvus_client.prepare_index_params()
|
||||
index_params.add_index(
|
||||
field_name="vector",
|
||||
index_type="HNSW",
|
||||
metric_type="COSINE",
|
||||
params={"M": 16, "efConstruction": 256}
|
||||
)
|
||||
milvus_client.create_index(collection_name=COLLECTION_NAME, index_params=index_params)
|
||||
print("成功为向量字段创建 HNSW 索引。")
|
||||
milvus_client.load_collection(collection_name=COLLECTION_NAME)
|
||||
print("已加载 Collection 到内存中。")
|
||||
|
||||
# 6. 执行多模态检索
|
||||
print(f"\n--> 正在 '{COLLECTION_NAME}' 中执行多模态检索")
|
||||
|
||||
# 示例1：图像+文本查询
|
||||
query_image_path = os.path.join(DATA_DIR, "query.png")
|
||||
query_text = "悬崖上的巨龙"
|
||||
query_vector = encoder.encode_query(image_path=query_image_path, text=query_text)
|
||||
|
||||
print(f"\n=== 多模态查询（图像+文本）===")
|
||||
print(f"查询图像: {query_image_path}")
|
||||
print(f"查询文本: {query_text}")
|
||||
|
||||
search_results = milvus_client.search(
|
||||
collection_name=COLLECTION_NAME,
|
||||
data=[query_vector],
|
||||
output_fields=["img_id", "image_path", "title", "description", "category", "location", "environment"],
|
||||
limit=6,
|
||||
search_params={"metric_type": "COSINE", "params": {"ef": 128}}
|
||||
)[0]
|
||||
|
||||
retrieved_results = []
|
||||
print("检索结果:")
|
||||
for i, hit in enumerate(search_results):
|
||||
print(f" Top {i+1}: ID={hit['id']}, 距离={hit['distance']:.4f}")
|
||||
print(f" 标题: {hit['entity']['title']}")
|
||||
print(f" 描述: {hit['entity']['description'][:100]}...")
|
||||
print(f" 类别: {hit['entity']['category']}")
|
||||
print(f" 路径: {hit['entity']['image_path']}")
|
||||
print("-" * 50)
|
||||
retrieved_results.append({
|
||||
'image_path': hit['entity']['image_path'],
|
||||
'distance': hit['distance']
|
||||
})
|
||||
|
||||
# 示例2：纯文本查询
|
||||
print(f"\n=== 纯文本查询 ===")
|
||||
text_query = "悬崖上的巨龙"
|
||||
text_query_vector = encoder.encode_query(text=text_query)
|
||||
|
||||
print(f"查询文本: {text_query}")
|
||||
|
||||
text_search_results = milvus_client.search(
|
||||
collection_name=COLLECTION_NAME,
|
||||
data=[text_query_vector],
|
||||
output_fields=["img_id", "image_path", "title", "description", "category", "location", "environment"],
|
||||
limit=3,
|
||||
search_params={"metric_type": "COSINE", "params": {"ef": 128}}
|
||||
)[0]
|
||||
|
||||
print("文本检索结果:")
|
||||
for i, hit in enumerate(text_search_results):
|
||||
print(f" Top {i+1}: {hit['entity']['title']} (距离: {hit['distance']:.4f})")
|
||||
print(f" 描述: {hit['entity']['description'][:80]}...")
|
||||
|
||||
# 示例3：纯图像查询
|
||||
print(f"\n=== 纯图像查询 ===")
|
||||
image_query_path = os.path.join(DATA_DIR, "query.png")
|
||||
image_query_vector = encoder.encode_query(image_path=image_query_path)
|
||||
|
||||
print(f"查询图像: {image_query_path}")
|
||||
|
||||
image_search_results = milvus_client.search(
|
||||
collection_name=COLLECTION_NAME,
|
||||
data=[image_query_vector],
|
||||
output_fields=["img_id", "image_path", "title", "description", "category", "location", "environment"],
|
||||
limit=3,
|
||||
search_params={"metric_type": "COSINE", "params": {"ef": 128}}
|
||||
)[0]
|
||||
|
||||
print("图像检索结果:")
|
||||
for i, hit in enumerate(image_search_results):
|
||||
print(f" Top {i+1}: {hit['entity']['title']} (距离: {hit['distance']:.4f})")
|
||||
print(f" 类别: {hit['entity']['category']}")
|
||||
print(f" 描述: {hit['entity']['description'][:80]}...")
|
||||
print(f" 路径: {hit['entity']['image_path']}")
|
||||
print("-" * 30)
|
||||
|
||||
# 7. 可视化与清理
|
||||
print(f"\n--> 正在可视化结果并清理资源")
|
||||
if retrieved_results:
|
||||
panoramic_image = visualize_results(query_image_path, retrieved_results)
|
||||
combined_image_path = "../../data/C4/multimodal_search.png"
|
||||
cv2.imwrite(combined_image_path, panoramic_image)
|
||||
print(f"结果图像已保存到: {combined_image_path}")
|
||||
# Image.open(combined_image_path).show()
|
||||
|
||||
# 8. 清理资源
|
||||
milvus_client.release_collection(collection_name=COLLECTION_NAME)
|
||||
print(f"已从内存中释放 Collection: '{COLLECTION_NAME}'")
|
||||
milvus_client.drop_collection(COLLECTION_NAME)
|
||||
print(f"已删除 Collection: '{COLLECTION_NAME}'")
|
||||
@@ -0,0 +1,209 @@
|
||||
import json
|
||||
import os
|
||||
import numpy as np
|
||||
from pymilvus import connections, MilvusClient, FieldSchema, CollectionSchema, DataType, Collection, AnnSearchRequest, RRFRanker
|
||||
from pymilvus.model.hybrid import BGEM3EmbeddingFunction
|
||||
|
||||
# 1. 初始化设置
|
||||
COLLECTION_NAME = "dragon_hybrid_demo"
|
||||
MILVUS_URI = "http://localhost:19530" # 服务器模式
|
||||
DATA_PATH = "../../data/C4/metadata/dragon.json" # 相对路径
|
||||
BATCH_SIZE = 50
|
||||
|
||||
# 2. 连接 Milvus 并初始化嵌入模型
|
||||
print(f"--> 正在连接到 Milvus: {MILVUS_URI}")
|
||||
connections.connect(uri=MILVUS_URI)
|
||||
|
||||
print("--> 正在初始化 BGE-M3 嵌入模型...")
|
||||
ef = BGEM3EmbeddingFunction(use_fp16=False, device="cpu")
|
||||
print(f"--> 嵌入模型初始化完成。密集向量维度: {ef.dim['dense']}")
|
||||
|
||||
# 3. 创建 Collection
|
||||
milvus_client = MilvusClient(uri=MILVUS_URI)
|
||||
if milvus_client.has_collection(COLLECTION_NAME):
|
||||
print(f"--> 正在删除已存在的 Collection '{COLLECTION_NAME}'...")
|
||||
milvus_client.drop_collection(COLLECTION_NAME)
|
||||
|
||||
fields = [
|
||||
FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100),
|
||||
FieldSchema(name="img_id", dtype=DataType.VARCHAR, max_length=100),
|
||||
FieldSchema(name="path", dtype=DataType.VARCHAR, max_length=256),
|
||||
FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
|
||||
FieldSchema(name="description", dtype=DataType.VARCHAR, max_length=4096),
|
||||
FieldSchema(name="category", dtype=DataType.VARCHAR, max_length=64),
|
||||
FieldSchema(name="location", dtype=DataType.VARCHAR, max_length=128),
|
||||
FieldSchema(name="environment", dtype=DataType.VARCHAR, max_length=64),
|
||||
FieldSchema(name="sparse_vector", dtype=DataType.SPARSE_FLOAT_VECTOR),
|
||||
FieldSchema(name="dense_vector", dtype=DataType.FLOAT_VECTOR, dim=ef.dim["dense"])
|
||||
]
|
||||
|
||||
# 如果集合不存在，则创建它及索引
|
||||
if not milvus_client.has_collection(COLLECTION_NAME):
|
||||
print(f"--> 正在创建 Collection '{COLLECTION_NAME}'...")
|
||||
schema = CollectionSchema(fields, description="关于龙的混合检索示例")
|
||||
# 创建集合
|
||||
collection = Collection(name=COLLECTION_NAME, schema=schema, consistency_level="Strong")
|
||||
print("--> Collection 创建成功。")
|
||||
|
||||
# 4. 创建索引
|
||||
print("--> 正在为新集合创建索引...")
|
||||
sparse_index = {"index_type": "SPARSE_INVERTED_INDEX", "metric_type": "IP"}
|
||||
collection.create_index("sparse_vector", sparse_index)
|
||||
print("稀疏向量索引创建成功。")
|
||||
|
||||
dense_index = {"index_type": "AUTOINDEX", "metric_type": "IP"}
|
||||
collection.create_index("dense_vector", dense_index)
|
||||
print("密集向量索引创建成功。")
|
||||
|
||||
collection = Collection(COLLECTION_NAME)
|
||||
|
||||
# 5. 加载数据并插入
|
||||
collection.load()
|
||||
print(f"--> Collection '{COLLECTION_NAME}' 已加载到内存。")
|
||||
|
||||
if collection.is_empty:
|
||||
print(f"--> Collection 为空，开始插入数据...")
|
||||
if not os.path.exists(DATA_PATH):
|
||||
raise FileNotFoundError(f"数据文件未找到: {DATA_PATH}")
|
||||
with open(DATA_PATH, 'r', encoding='utf-8') as f:
|
||||
dataset = json.load(f)
|
||||
|
||||
docs, metadata = [], []
|
||||
for item in dataset:
|
||||
parts = [
|
||||
item.get('title', ''),
|
||||
item.get('description', ''),
|
||||
item.get('location', ''),
|
||||
item.get('environment', ''),
|
||||
# *item.get('combat_details', {}).get('combat_style', []),
|
||||
# *item.get('combat_details', {}).get('abilities_used', []),
|
||||
# item.get('scene_info', {}).get('time_of_day', '')
|
||||
]
|
||||
docs.append(' '.join(filter(None, parts)))
|
||||
metadata.append(item)
|
||||
print(f"--> 数据加载完成，共 {len(docs)} 条。")
|
||||
|
||||
print("--> 正在生成向量嵌入...")
|
||||
embeddings = ef(docs)
|
||||
print("--> 向量生成完成。")
|
||||
|
||||
print("--> 正在分批插入数据...")
|
||||
# 为每个字段准备批量数据
|
||||
img_ids = [doc["img_id"] for doc in metadata]
|
||||
paths = [doc["path"] for doc in metadata]
|
||||
titles = [doc["title"] for doc in metadata]
|
||||
descriptions = [doc["description"] for doc in metadata]
|
||||
categories = [doc["category"] for doc in metadata]
|
||||
locations = [doc["location"] for doc in metadata]
|
||||
environments = [doc["environment"] for doc in metadata]
|
||||
|
||||
# 获取向量
|
||||
sparse_vectors = embeddings["sparse"]
|
||||
dense_vectors = embeddings["dense"]
|
||||
|
||||
# 插入数据
|
||||
collection.insert([
|
||||
img_ids,
|
||||
paths,
|
||||
titles,
|
||||
descriptions,
|
||||
categories,
|
||||
locations,
|
||||
environments,
|
||||
sparse_vectors,
|
||||
dense_vectors
|
||||
])
|
||||
|
||||
collection.flush()
|
||||
print(f"--> 数据插入完成，总数: {collection.num_entities}")
|
||||
else:
|
||||
print(f"--> Collection 中已有 {collection.num_entities} 条数据，跳过插入。")
|
||||
|
||||
# 6. 执行搜索
|
||||
search_query = "悬崖上的巨龙"
|
||||
search_filter = 'category in ["western_dragon", "chinese_dragon", "movie_character"]'
|
||||
top_k = 5
|
||||
|
||||
print(f"\n{'='*20} 开始混合搜索 {'='*20}")
|
||||
print(f"查询: '{search_query}'")
|
||||
print(f"过滤器: '{search_filter}'")
|
||||
|
||||
query_embeddings = ef([search_query])
|
||||
dense_vec = query_embeddings["dense"][0]
|
||||
sparse_vec = query_embeddings["sparse"]._getrow(0)
|
||||
|
||||
# 打印向量信息
|
||||
print("\n=== 向量信息 ===")
|
||||
print(f"密集向量维度: {len(dense_vec)}")
|
||||
print(f"密集向量前5个元素: {dense_vec[:5]}")
|
||||
print(f"密集向量范数: {np.linalg.norm(dense_vec):.4f}")
|
||||
|
||||
print(f"\n稀疏向量维度: {sparse_vec.shape[1]}")
|
||||
print(f"稀疏向量非零元素数量: {sparse_vec.nnz}")
|
||||
print("稀疏向量前5个非零元素:")
|
||||
for i in range(min(5, sparse_vec.nnz)):
|
||||
print(f" - 索引: {sparse_vec.indices[i]}, 值: {sparse_vec.data[i]:.4f}")
|
||||
density = (sparse_vec.nnz / sparse_vec.shape[1] * 100)
|
||||
print(f"\n稀疏向量密度: {density:.8f}%")
|
||||
|
||||
# 定义搜索参数
|
||||
search_params = {"metric_type": "IP", "params": {}}
|
||||
|
||||
# 先执行单独的搜索
|
||||
print("\n--- [单独] 密集向量搜索结果 ---")
|
||||
dense_results = collection.search(
|
||||
[dense_vec],
|
||||
anns_field="dense_vector",
|
||||
param=search_params,
|
||||
limit=top_k,
|
||||
expr=search_filter,
|
||||
output_fields=["title", "path", "description", "category", "location", "environment"]
|
||||
)[0]
|
||||
|
||||
for i, hit in enumerate(dense_results):
|
||||
print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
|
||||
print(f" 路径: {hit.entity.get('path')}")
|
||||
print(f" 描述: {hit.entity.get('description')[:100]}...")
|
||||
|
||||
print("\n--- [单独] 稀疏向量搜索结果 ---")
|
||||
sparse_results = collection.search(
|
||||
[sparse_vec],
|
||||
anns_field="sparse_vector",
|
||||
param=search_params,
|
||||
limit=top_k,
|
||||
expr=search_filter,
|
||||
output_fields=["title", "path", "description", "category", "location", "environment"]
|
||||
)[0]
|
||||
|
||||
for i, hit in enumerate(sparse_results):
|
||||
print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
|
||||
print(f" 路径: {hit.entity.get('path')}")
|
||||
print(f" 描述: {hit.entity.get('description')[:100]}...")
|
||||
|
||||
print("\n--- [混合] 稀疏+密集向量搜索结果 ---")
|
||||
# 创建 RRF 融合器
|
||||
rerank = RRFRanker(k=60)
|
||||
|
||||
# 创建搜索请求
|
||||
dense_req = AnnSearchRequest([dense_vec], "dense_vector", search_params, limit=top_k)
|
||||
sparse_req = AnnSearchRequest([sparse_vec], "sparse_vector", search_params, limit=top_k)
|
||||
|
||||
# 执行混合搜索
|
||||
results = collection.hybrid_search(
|
||||
[sparse_req, dense_req],
|
||||
rerank=rerank,
|
||||
limit=top_k,
|
||||
output_fields=["title", "path", "description", "category", "location", "environment"]
|
||||
)[0]
|
||||
|
||||
# 打印最终结果
|
||||
for i, hit in enumerate(results):
|
||||
print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
|
||||
print(f" 路径: {hit.entity.get('path')}")
|
||||
print(f" 描述: {hit.entity.get('description')[:100]}...")
|
||||
|
||||
# 7. 清理资源
|
||||
milvus_client.release_collection(collection_name=COLLECTION_NAME)
|
||||
print(f"已从内存中释放 Collection: '{COLLECTION_NAME}'")
|
||||
milvus_client.drop_collection(COLLECTION_NAME)
|
||||
print(f"已删除 Collection: '{COLLECTION_NAME}'")
|
||||
@@ -0,0 +1,329 @@
|
||||
import json
|
||||
import os
|
||||
import numpy as np
|
||||
import torch
|
||||
from transformers import AutoModel, AutoProcessor
|
||||
from sklearn.feature_extraction.text import TfidfVectorizer
|
||||
from scipy.sparse import csr_matrix
|
||||
from pymilvus import connections, MilvusClient, FieldSchema, CollectionSchema, DataType, Collection, AnnSearchRequest, RRFRanker
|
||||
|
||||
# 1. 初始化设置
|
||||
COLLECTION_NAME = "dragon_siglip_demo"
|
||||
MILVUS_URI = "http://localhost:19530" # 服务器模式
|
||||
DATA_PATH = "../../data/C4/metadata/dragon.json" # 相对路径
|
||||
BATCH_SIZE = 50
|
||||
|
||||
# 2. 自定义SigLIP嵌入函数类
|
||||
class SigLIPEmbeddingFunction:
|
||||
def __init__(self, model_name="google/siglip-base-patch16-256-multilingual", device="cpu"):
|
||||
"""
|
||||
初始化SigLIP嵌入函数
|
||||
Args:
|
||||
model_name: SigLIP模型名称
|
||||
device: 设备类型 ("cpu" 或 "cuda")
|
||||
"""
|
||||
self.model_name = model_name
|
||||
self.device = device
|
||||
|
||||
print(f"--> 正在加载 SigLIP 模型: {model_name}")
|
||||
self.model = AutoModel.from_pretrained(model_name)
|
||||
self.processor = AutoProcessor.from_pretrained(model_name)
|
||||
self.model.to(device)
|
||||
self.model.eval()
|
||||
|
||||
# 初始化TF-IDF作为稀疏向量生成器
|
||||
self.tfidf_vectorizer = TfidfVectorizer(
|
||||
max_features=10000, # 限制词汇表大小以节省空间
|
||||
stop_words='english',
|
||||
ngram_range=(1, 2)
|
||||
)
|
||||
self.tfidf_fitted = False
|
||||
|
||||
# 获取文本编码器的输出维度
|
||||
with torch.no_grad():
|
||||
dummy_text = ["test"]
|
||||
inputs = self.processor(text=dummy_text, padding="max_length", return_tensors="pt")
|
||||
outputs = self.model.text_model(**{k: v.to(device) for k, v in inputs.items() if k != 'pixel_values'})
|
||||
self.dense_dim = outputs.pooler_output.shape[-1]
|
||||
|
||||
print(f"--> SigLIP 模型加载完成。密集向量维度: {self.dense_dim}")
|
||||
|
||||
@property
|
||||
def dim(self):
|
||||
"""返回维度信息，兼容原BGE-M3接口"""
|
||||
return {
|
||||
"dense": self.dense_dim,
|
||||
"sparse": self.tfidf_vectorizer.max_features if self.tfidf_fitted else 10000
|
||||
}
|
||||
|
||||
def fit_sparse(self, docs):
|
||||
"""拟合稀疏向量模型（TF-IDF）"""
|
||||
print("--> 正在拟合 TF-IDF 模型...")
|
||||
self.tfidf_vectorizer.fit(docs)
|
||||
self.tfidf_fitted = True
|
||||
print(f"--> TF-IDF 模型拟合完成。词汇表大小: {len(self.tfidf_vectorizer.vocabulary_)}")
|
||||
|
||||
def encode_text_dense(self, texts):
|
||||
"""使用SigLIP编码文本为密集向量"""
|
||||
if isinstance(texts, str):
|
||||
texts = [texts]
|
||||
|
||||
dense_vectors = []
|
||||
batch_size = 8 # 减小批次大小以节省内存
|
||||
|
||||
with torch.no_grad():
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch_texts = texts[i:i + batch_size]
|
||||
inputs = self.processor(text=batch_texts, padding="max_length", truncation=True, return_tensors="pt")
|
||||
inputs = {k: v.to(self.device) for k, v in inputs.items() if k != 'pixel_values'}
|
||||
|
||||
outputs = self.model.text_model(**inputs)
|
||||
embeddings = outputs.pooler_output
|
||||
|
||||
# 归一化向量
|
||||
embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
|
||||
dense_vectors.extend(embeddings.cpu().numpy())
|
||||
|
||||
return np.array(dense_vectors)
|
||||
|
||||
def encode_text_sparse(self, texts):
|
||||
"""使用TF-IDF编码文本为稀疏向量"""
|
||||
if not self.tfidf_fitted:
|
||||
raise ValueError("请先调用 fit_sparse() 方法拟合TF-IDF模型")
|
||||
|
||||
if isinstance(texts, str):
|
||||
texts = [texts]
|
||||
|
||||
sparse_matrix = self.tfidf_vectorizer.transform(texts)
|
||||
return sparse_matrix
|
||||
|
||||
def __call__(self, texts):
|
||||
"""主调用方法，返回密集和稀疏向量"""
|
||||
if isinstance(texts, str):
|
||||
texts = [texts]
|
||||
|
||||
# 如果还没有拟合稀疏模型，先拟合
|
||||
if not self.tfidf_fitted:
|
||||
self.fit_sparse(texts)
|
||||
|
||||
dense_vectors = self.encode_text_dense(texts)
|
||||
sparse_vectors = self.encode_text_sparse(texts)
|
||||
|
||||
return {
|
||||
"dense": dense_vectors,
|
||||
"sparse": sparse_vectors
|
||||
}
|
||||
|
||||
# 3. 连接 Milvus 并初始化嵌入模型
|
||||
print(f"--> 正在连接到 Milvus: {MILVUS_URI}")
|
||||
connections.connect(uri=MILVUS_URI)
|
||||
|
||||
print("--> 正在初始化 SigLIP 嵌入模型...")
|
||||
ef = SigLIPEmbeddingFunction(device="cpu") # 如果有GPU可以改为"cuda"
|
||||
|
||||
# 4. 创建 Collection
|
||||
milvus_client = MilvusClient(uri=MILVUS_URI)
|
||||
if milvus_client.has_collection(COLLECTION_NAME):
|
||||
print(f"--> 正在删除已存在的 Collection '{COLLECTION_NAME}'...")
|
||||
milvus_client.drop_collection(COLLECTION_NAME)
|
||||
|
||||
fields = [
|
||||
FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100),
|
||||
FieldSchema(name="img_id", dtype=DataType.VARCHAR, max_length=100),
|
||||
FieldSchema(name="path", dtype=DataType.VARCHAR, max_length=256),
|
||||
FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
|
||||
FieldSchema(name="description", dtype=DataType.VARCHAR, max_length=4096),
|
||||
FieldSchema(name="category", dtype=DataType.VARCHAR, max_length=64),
|
||||
FieldSchema(name="location", dtype=DataType.VARCHAR, max_length=128),
|
||||
FieldSchema(name="environment", dtype=DataType.VARCHAR, max_length=64),
|
||||
FieldSchema(name="sparse_vector", dtype=DataType.SPARSE_FLOAT_VECTOR),
|
||||
FieldSchema(name="dense_vector", dtype=DataType.FLOAT_VECTOR, dim=ef.dim["dense"])
|
||||
]
|
||||
|
||||
# 如果集合不存在，则创建它及索引
|
||||
if not milvus_client.has_collection(COLLECTION_NAME):
|
||||
print(f"--> 正在创建 Collection '{COLLECTION_NAME}'...")
|
||||
schema = CollectionSchema(fields, description="使用SigLIP的龙混合检索示例")
|
||||
# 创建集合
|
||||
collection = Collection(name=COLLECTION_NAME, schema=schema, consistency_level="Strong")
|
||||
print("--> Collection 创建成功。")
|
||||
|
||||
# 5. 创建索引
|
||||
print("--> 正在为新集合创建索引...")
|
||||
sparse_index = {"index_type": "SPARSE_INVERTED_INDEX", "metric_type": "IP"}
|
||||
collection.create_index("sparse_vector", sparse_index)
|
||||
print("稀疏向量索引创建成功。")
|
||||
|
||||
dense_index = {"index_type": "AUTOINDEX", "metric_type": "IP"}
|
||||
collection.create_index("dense_vector", dense_index)
|
||||
print("密集向量索引创建成功。")
|
||||
|
||||
collection = Collection(COLLECTION_NAME)
|
||||
|
||||
# 6. 加载数据并插入
|
||||
collection.load()
|
||||
print(f"--> Collection '{COLLECTION_NAME}' 已加载到内存。")
|
||||
|
||||
if collection.is_empty:
|
||||
print(f"--> Collection 为空，开始插入数据...")
|
||||
if not os.path.exists(DATA_PATH):
|
||||
raise FileNotFoundError(f"数据文件未找到: {DATA_PATH}")
|
||||
with open(DATA_PATH, 'r', encoding='utf-8') as f:
|
||||
dataset = json.load(f)
|
||||
|
||||
docs, metadata = [], []
|
||||
for item in dataset:
|
||||
parts = [
|
||||
item.get('title', ''),
|
||||
item.get('description', ''),
|
||||
item.get('location', ''),
|
||||
item.get('environment', ''),
|
||||
# *item.get('combat_details', {}).get('combat_style', []),
|
||||
# *item.get('combat_details', {}).get('abilities_used', []),
|
||||
# item.get('scene_info', {}).get('time_of_day', '')
|
||||
]
|
||||
docs.append(' '.join(filter(None, parts)))
|
||||
metadata.append(item)
|
||||
print(f"--> 数据加载完成，共 {len(docs)} 条。")
|
||||
|
||||
print("--> 正在生成向量嵌入...")
|
||||
embeddings = ef(docs)
|
||||
print("--> 向量生成完成。")
|
||||
|
||||
print("--> 正在分批插入数据...")
|
||||
# 为每个字段准备批量数据
|
||||
img_ids = [doc["img_id"] for doc in metadata]
|
||||
paths = [doc["path"] for doc in metadata]
|
||||
titles = [doc["title"] for doc in metadata]
|
||||
descriptions = [doc["description"] for doc in metadata]
|
||||
categories = [doc["category"] for doc in metadata]
|
||||
locations = [doc["location"] for doc in metadata]
|
||||
environments = [doc["environment"] for doc in metadata]
|
||||
|
||||
# 获取向量 - 注意SigLIP返回的格式与BGE-M3不同
|
||||
sparse_vectors = []
|
||||
dense_vectors = embeddings["dense"].tolist()
|
||||
|
||||
# 将稀疏矩阵转换为Milvus可接受的格式
|
||||
sparse_matrix = embeddings["sparse"]
|
||||
for i in range(sparse_matrix.shape[0]):
|
||||
row = sparse_matrix.getrow(i)
|
||||
# 创建稀疏向量字典格式
|
||||
sparse_dict = {}
|
||||
for j in range(row.nnz):
|
||||
sparse_dict[row.indices[j]] = float(row.data[j])
|
||||
sparse_vectors.append(sparse_dict)
|
||||
|
||||
# 插入数据
|
||||
collection.insert([
|
||||
img_ids,
|
||||
paths,
|
||||
titles,
|
||||
descriptions,
|
||||
categories,
|
||||
locations,
|
||||
environments,
|
||||
sparse_vectors,
|
||||
dense_vectors
|
||||
])
|
||||
|
||||
collection.flush()
|
||||
print(f"--> 数据插入完成，总数: {collection.num_entities}")
|
||||
else:
|
||||
print(f"--> Collection 中已有 {collection.num_entities} 条数据，跳过插入。")
|
||||
|
||||
# 7. 执行搜索
|
||||
search_query = "悬崖上的巨龙"
|
||||
search_filter = 'category in ["western_dragon", "chinese_dragon", "movie_character"]'
|
||||
top_k = 5
|
||||
|
||||
print(f"\n{'='*20} 开始混合搜索 {'='*20}")
|
||||
print(f"查询: '{search_query}'")
|
||||
print(f"过滤器: '{search_filter}'")
|
||||
|
||||
# 生成查询向量
|
||||
query_embeddings = ef([search_query])
|
||||
dense_vec = query_embeddings["dense"][0].tolist()
|
||||
|
||||
# 处理稀疏向量
|
||||
sparse_matrix = query_embeddings["sparse"]
|
||||
sparse_row = sparse_matrix.getrow(0)
|
||||
sparse_dict = {}
|
||||
for j in range(sparse_row.nnz):
|
||||
sparse_dict[sparse_row.indices[j]] = float(sparse_row.data[j])
|
||||
|
||||
# 打印向量信息
|
||||
print("\n=== 向量信息 ===")
|
||||
print(f"密集向量维度: {len(dense_vec)}")
|
||||
print(f"密集向量前5个元素: {dense_vec[:5]}")
|
||||
print(f"密集向量范数: {np.linalg.norm(dense_vec):.4f}")
|
||||
|
||||
print(f"\n稀疏向量维度: {sparse_matrix.shape[1]}")
|
||||
print(f"稀疏向量非零元素数量: {sparse_row.nnz}")
|
||||
print("稀疏向量前5个非零元素:")
|
||||
for i, (idx, val) in enumerate(list(sparse_dict.items())[:5]):
|
||||
print(f" - 索引: {idx}, 值: {val:.4f}")
|
||||
density = (sparse_row.nnz / sparse_matrix.shape[1] * 100)
|
||||
print(f"\n稀疏向量密度: {density:.8f}%")
|
||||
|
||||
# 定义搜索参数
|
||||
search_params = {"metric_type": "IP", "params": {}}
|
||||
|
||||
# 先执行单独的搜索
|
||||
print("\n--- [单独] 密集向量搜索结果 ---")
|
||||
dense_results = collection.search(
|
||||
[dense_vec],
|
||||
anns_field="dense_vector",
|
||||
param=search_params,
|
||||
limit=top_k,
|
||||
expr=search_filter,
|
||||
output_fields=["title", "path", "description", "category", "location", "environment"]
|
||||
)[0]
|
||||
|
||||
for i, hit in enumerate(dense_results):
|
||||
print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
|
||||
print(f" 路径: {hit.entity.get('path')}")
|
||||
print(f" 描述: {hit.entity.get('description')[:100]}...")
|
||||
|
||||
print("\n--- [单独] 稀疏向量搜索结果 ---")
|
||||
sparse_results = collection.search(
|
||||
[sparse_dict],
|
||||
anns_field="sparse_vector",
|
||||
param=search_params,
|
||||
limit=top_k,
|
||||
expr=search_filter,
|
||||
output_fields=["title", "path", "description", "category", "location", "environment"]
|
||||
)[0]
|
||||
|
||||
for i, hit in enumerate(sparse_results):
|
||||
print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
|
||||
print(f" 路径: {hit.entity.get('path')}")
|
||||
print(f" 描述: {hit.entity.get('description')[:100]}...")
|
||||
|
||||
print("\n--- [混合] 稀疏+密集向量搜索结果 ---")
|
||||
# 创建 RRF 融合器
|
||||
rerank = RRFRanker(k=60)
|
||||
|
||||
# 创建搜索请求
|
||||
dense_req = AnnSearchRequest([dense_vec], "dense_vector", search_params, limit=top_k)
|
||||
sparse_req = AnnSearchRequest([sparse_dict], "sparse_vector", search_params, limit=top_k)
|
||||
|
||||
# 执行混合搜索
|
||||
results = collection.hybrid_search(
|
||||
[sparse_req, dense_req],
|
||||
rerank=rerank,
|
||||
limit=top_k,
|
||||
output_fields=["title", "path", "description", "category", "location", "environment"]
|
||||
)[0]
|
||||
|
||||
# 打印最终结果
|
||||
for i, hit in enumerate(results):
|
||||
print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
|
||||
print(f" 路径: {hit.entity.get('path')}")
|
||||
print(f" 描述: {hit.entity.get('description')[:100]}...")
|
||||
|
||||
# 8. 清理资源
|
||||
milvus_client.release_collection(collection_name=COLLECTION_NAME)
|
||||
print(f"已从内存中释放 Collection: '{COLLECTION_NAME}'")
|
||||
milvus_client.drop_collection(COLLECTION_NAME)
|
||||
print(f"已删除 Collection: '{COLLECTION_NAME}'")
|
||||
@@ -0,0 +1,111 @@
|
||||
import os
|
||||
from langchain_deepseek import ChatDeepSeek
|
||||
from langchain_community.document_loaders import BiliBiliLoader
|
||||
from langchain.chains.query_constructor.base import AttributeInfo
|
||||
from langchain.retrievers.self_query.base import SelfQueryRetriever
|
||||
from langchain_community.vectorstores import Chroma
|
||||
from langchain_huggingface import HuggingFaceEmbeddings
|
||||
import logging
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
# 1. 初始化视频数据
|
||||
video_urls = [
|
||||
"https://www.bilibili.com/video/BV1Bo4y1A7FU",
|
||||
"https://www.bilibili.com/video/BV1ug4y157xA",
|
||||
"https://www.bilibili.com/video/BV1yh411V7ge",
|
||||
]
|
||||
|
||||
bili = []
|
||||
try:
|
||||
loader = BiliBiliLoader(video_urls=video_urls)
|
||||
docs = loader.load()
|
||||
|
||||
for doc in docs:
|
||||
original = doc.metadata
|
||||
|
||||
# 提取基本元数据字段
|
||||
metadata = {
|
||||
'title': original.get('title', '未知标题'),
|
||||
'author': original.get('owner', {}).get('name', '未知作者'),
|
||||
'source': original.get('bvid', '未知ID'),
|
||||
'view_count': original.get('stat', {}).get('view', 0),
|
||||
'length': original.get('duration', 0),
|
||||
}
|
||||
|
||||
doc.metadata = metadata
|
||||
bili.append(doc)
|
||||
|
||||
except Exception as e:
|
||||
print(f"加载BiliBili视频失败: {str(e)}")
|
||||
|
||||
if not bili:
|
||||
print("没有成功加载任何视频，程序退出")
|
||||
exit()
|
||||
|
||||
# 2. 创建向量存储
|
||||
embed_model = HuggingFaceEmbeddings(model_name="BAAI/bge-small-zh-v1.5")
|
||||
vectorstore = Chroma.from_documents(bili, embed_model)
|
||||
|
||||
# 3. 配置元数据字段信息
|
||||
metadata_field_info = [
|
||||
AttributeInfo(
|
||||
name="title",
|
||||
description="视频标题（字符串）",
|
||||
type="string",
|
||||
),
|
||||
AttributeInfo(
|
||||
name="author",
|
||||
description="视频作者（字符串）",
|
||||
type="string",
|
||||
),
|
||||
AttributeInfo(
|
||||
name="view_count",
|
||||
description="视频观看次数（整数）",
|
||||
type="integer",
|
||||
),
|
||||
AttributeInfo(
|
||||
name="length",
|
||||
description="视频长度（整数）",
|
||||
type="integer"
|
||||
)
|
||||
]
|
||||
|
||||
# 4. 创建自查询检索器
|
||||
llm = ChatDeepSeek(
|
||||
model="deepseek-chat",
|
||||
temperature=0,
|
||||
api_key=os.getenv("DEEPSEEK_API_KEY")
|
||||
)
|
||||
|
||||
retriever = SelfQueryRetriever.from_llm(
|
||||
llm=llm,
|
||||
vectorstore=vectorstore,
|
||||
document_contents="记录视频标题、作者、观看次数等信息的视频元数据",
|
||||
metadata_field_info=metadata_field_info,
|
||||
enable_limit=True,
|
||||
verbose=True
|
||||
)
|
||||
|
||||
# 5. 执行查询示例
|
||||
queries = [
|
||||
"时间最短的视频",
|
||||
"时长大于600秒的视频"
|
||||
]
|
||||
|
||||
for query in queries:
|
||||
print(f"\n--- 查询: '{query}' ---")
|
||||
results = retriever.invoke(query)
|
||||
if results:
|
||||
for doc in results:
|
||||
title = doc.metadata.get('title', '未知标题')
|
||||
author = doc.metadata.get('author', '未知作者')
|
||||
view_count = doc.metadata.get('view_count', '未知')
|
||||
length = doc.metadata.get('length', '未知')
|
||||
print(f"标题: {title}")
|
||||
print(f"作者: {author}")
|
||||
print(f"观看次数: {view_count}")
|
||||
print(f"时长: {length}秒")
|
||||
print("="*50)
|
||||
else:
|
||||
print("未找到匹配的视频")
|
||||
@@ -0,0 +1,220 @@
|
||||
import os
|
||||
import sys
|
||||
import sqlite3
|
||||
|
||||
# 添加text2sql模块路径
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), 'text2sql'))
|
||||
|
||||
from text2sql.text2sql_agent import SimpleText2SQLAgent
|
||||
|
||||
|
||||
def setup_demo():
|
||||
"""设置演示环境"""
|
||||
print("=== Text2SQL框架演示 ===\n")
|
||||
|
||||
# 检查API密钥
|
||||
api_key = os.getenv("DEEPSEEK_API_KEY")
|
||||
if not api_key:
|
||||
print("先设置DEEPSEEK_API_KEY环境变量")
|
||||
return None
|
||||
|
||||
# 创建演示数据库
|
||||
print("创建演示数据库...")
|
||||
db_path = create_demo_database()
|
||||
|
||||
# 初始化Text2SQL代理
|
||||
print("初始化Text2SQL代理...")
|
||||
agent = SimpleText2SQLAgent(api_key=api_key)
|
||||
|
||||
# 连接数据库
|
||||
print("连接数据库...")
|
||||
if not agent.connect_database(db_path):
|
||||
print("数据库连接失败!")
|
||||
return None
|
||||
|
||||
# 加载知识库
|
||||
print("加载知识库...")
|
||||
try:
|
||||
agent.load_knowledge_base()
|
||||
print("知识库加载成功!")
|
||||
except Exception as e:
|
||||
print(f"知识库加载失败: {str(e)}")
|
||||
return None
|
||||
|
||||
return agent, db_path
|
||||
|
||||
|
||||
def create_demo_database():
|
||||
"""创建演示数据库"""
|
||||
db_path = "text2sql_demo.db"
|
||||
|
||||
if os.path.exists(db_path):
|
||||
os.remove(db_path)
|
||||
|
||||
conn = sqlite3.connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 创建用户表
|
||||
cursor.execute("""
|
||||
CREATE TABLE users (
|
||||
id INTEGER PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
email TEXT UNIQUE,
|
||||
age INTEGER,
|
||||
city TEXT
|
||||
)
|
||||
""")
|
||||
|
||||
# 创建产品表
|
||||
cursor.execute("""
|
||||
CREATE TABLE products (
|
||||
id INTEGER PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
category TEXT,
|
||||
price REAL,
|
||||
stock INTEGER
|
||||
)
|
||||
""")
|
||||
|
||||
# 创建订单表
|
||||
cursor.execute("""
|
||||
CREATE TABLE orders (
|
||||
id INTEGER PRIMARY KEY,
|
||||
user_id INTEGER,
|
||||
product_id INTEGER,
|
||||
quantity INTEGER,
|
||||
order_date TEXT,
|
||||
total_price REAL,
|
||||
FOREIGN KEY (user_id) REFERENCES users(id),
|
||||
FOREIGN KEY (product_id) REFERENCES products(id)
|
||||
)
|
||||
""")
|
||||
|
||||
# 插入示例数据
|
||||
users_data = [
|
||||
(1, '张三', 'zhangsan@email.com', 25, '北京'),
|
||||
(2, '李四', 'lisi@email.com', 32, '上海'),
|
||||
(3, '王五', 'wangwu@email.com', 28, '广州'),
|
||||
(4, '赵六', 'zhaoliu@email.com', 35, '深圳'),
|
||||
(5, '陈七', 'chenqi@email.com', 29, '杭州'),
|
||||
]
|
||||
|
||||
products_data = [
|
||||
(1, 'iPhone 15', '电子产品', 7999.0, 50),
|
||||
(2, 'MacBook Pro', '电子产品', 12999.0, 20),
|
||||
(3, 'Nike运动鞋', '服装', 599.0, 100),
|
||||
(4, '办公椅', '家具', 899.0, 30),
|
||||
(5, '台灯', '家具', 199.0, 80),
|
||||
(6, 'iPad', '电子产品', 3999.0, 40),
|
||||
(7, 'Adidas外套', '服装', 399.0, 60),
|
||||
]
|
||||
|
||||
orders_data = [
|
||||
(1, 1, 1, 1, '2024-01-15', 7999.0),
|
||||
(2, 2, 3, 2, '2024-01-16', 1198.0),
|
||||
(3, 3, 5, 1, '2024-01-17', 199.0),
|
||||
(4, 1, 2, 1, '2024-01-18', 12999.0),
|
||||
(5, 4, 4, 1, '2024-01-19', 899.0),
|
||||
(6, 5, 6, 1, '2024-01-20', 3999.0),
|
||||
(7, 2, 7, 1, '2024-01-21', 399.0),
|
||||
]
|
||||
|
||||
cursor.executemany("INSERT INTO users VALUES (?, ?, ?, ?, ?)", users_data)
|
||||
cursor.executemany("INSERT INTO products VALUES (?, ?, ?, ?, ?)", products_data)
|
||||
cursor.executemany("INSERT INTO orders VALUES (?, ?, ?, ?, ?, ?)", orders_data)
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
print(f"演示数据库已创建: {db_path}")
|
||||
return db_path
|
||||
|
||||
|
||||
def run_demo_queries(agent):
|
||||
"""运行演示查询"""
|
||||
demo_questions = [
|
||||
"查询所有用户的姓名和邮箱",
|
||||
"年龄大于30的用户有哪些",
|
||||
"哪些产品的库存少于50",
|
||||
"查询来自北京的用户的所有订单",
|
||||
"统计每个城市的用户数量",
|
||||
"查询价格在500-8000之间的产品"
|
||||
]
|
||||
|
||||
print("\n开始运行演示查询...\n")
|
||||
|
||||
success_count = 0
|
||||
|
||||
for i, question in enumerate(demo_questions, 1):
|
||||
print(f"问题 {i}: {question}")
|
||||
print("-" * 60)
|
||||
|
||||
try:
|
||||
result = agent.query(question)
|
||||
|
||||
if result["success"]:
|
||||
print(f"成功! SQL: {result['sql']}")
|
||||
|
||||
if isinstance(result["results"], dict) and "rows" in result["results"]:
|
||||
count = result["results"]["count"]
|
||||
print(f"返回 {count} 行数据")
|
||||
|
||||
# 显示前2行数据
|
||||
if count > 0:
|
||||
for j, row in enumerate(result["results"]["rows"][:2]):
|
||||
row_str = " | ".join(f"{k}: {v}" for k, v in row.items())
|
||||
print(f" {j+1}. {row_str}")
|
||||
|
||||
if count > 2:
|
||||
print(f" ... 还有 {count - 2} 行")
|
||||
else:
|
||||
print(f"结果: {result['results']}")
|
||||
|
||||
success_count += 1
|
||||
|
||||
else:
|
||||
print(f"失败: {result['error']}")
|
||||
print(f"SQL: {result['sql']}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"执行错误: {str(e)}")
|
||||
|
||||
print()
|
||||
|
||||
# 输出统计
|
||||
total_count = len(demo_questions)
|
||||
|
||||
|
||||
def cleanup(agent, db_path):
|
||||
"""清理资源"""
|
||||
print("\n清理资源...")
|
||||
|
||||
if agent:
|
||||
agent.cleanup()
|
||||
|
||||
if os.path.exists(db_path):
|
||||
os.remove(db_path)
|
||||
print(f"已删除演示数据库: {db_path}")
|
||||
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
# 设置演示环境
|
||||
setup_result = setup_demo()
|
||||
|
||||
if setup_result is None:
|
||||
return
|
||||
|
||||
agent, db_path = setup_result
|
||||
|
||||
try:
|
||||
# 运行演示查询
|
||||
run_demo_queries(agent)
|
||||
|
||||
finally:
|
||||
# 清理资源
|
||||
cleanup(agent, db_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,377 @@
|
||||
import os
|
||||
import json
|
||||
import sqlite3
|
||||
import numpy as np
|
||||
from typing import List, Dict, Any
|
||||
from sentence_transformers import SentenceTransformer
|
||||
from sklearn.feature_extraction.text import TfidfVectorizer
|
||||
from sklearn.metrics.pairwise import cosine_similarity
|
||||
from pymilvus import connections, MilvusClient, FieldSchema, CollectionSchema, DataType, Collection
|
||||
|
||||
|
||||
class BGESmallEmbeddingFunction:
|
||||
"""BGE-Small中文嵌入函数，用于Text2SQL知识库向量化"""
|
||||
|
||||
def __init__(self, model_name="BAAI/bge-small-zh-v1.5", device="cpu"):
|
||||
self.model_name = model_name
|
||||
self.device = device
|
||||
self.model = SentenceTransformer(model_name, device=device)
|
||||
self.dense_dim = self.model.get_sentence_embedding_dimension()
|
||||
|
||||
def encode_text(self, texts):
|
||||
"""编码文本为密集向量"""
|
||||
if isinstance(texts, str):
|
||||
texts = [texts]
|
||||
|
||||
embeddings = self.model.encode(
|
||||
texts,
|
||||
normalize_embeddings=True,
|
||||
batch_size=16,
|
||||
convert_to_numpy=True
|
||||
)
|
||||
|
||||
return embeddings
|
||||
|
||||
@property
|
||||
def dim(self):
|
||||
"""返回向量维度"""
|
||||
return self.dense_dim
|
||||
|
||||
|
||||
class SimpleKnowledgeBase:
|
||||
"""简化的知识库，使用BGE-Small进行向量检索"""
|
||||
|
||||
def __init__(self, milvus_uri: str = "http://localhost:19530"):
|
||||
self.milvus_uri = milvus_uri
|
||||
self.collection_name = "text2sql_knowledge_base"
|
||||
self.milvus_client = None
|
||||
self.collection = None
|
||||
|
||||
self.embedding_function = BGESmallEmbeddingFunction(
|
||||
model_name="BAAI/bge-small-zh-v1.5",
|
||||
device="cpu"
|
||||
)
|
||||
|
||||
self.sql_examples = []
|
||||
self.table_schemas = []
|
||||
self.data_loaded = False
|
||||
|
||||
def connect_milvus(self):
|
||||
"""连接Milvus数据库"""
|
||||
connections.connect(uri=self.milvus_uri)
|
||||
self.milvus_client = MilvusClient(uri=self.milvus_uri)
|
||||
return True
|
||||
|
||||
def create_collection(self):
|
||||
"""创建Milvus集合"""
|
||||
if not self.milvus_client:
|
||||
self.connect_milvus()
|
||||
|
||||
if self.milvus_client.has_collection(self.collection_name):
|
||||
self.milvus_client.drop_collection(self.collection_name)
|
||||
|
||||
fields = [
|
||||
FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100),
|
||||
FieldSchema(name="content_type", dtype=DataType.VARCHAR, max_length=50),
|
||||
FieldSchema(name="question", dtype=DataType.VARCHAR, max_length=1000),
|
||||
FieldSchema(name="sql", dtype=DataType.VARCHAR, max_length=2000),
|
||||
FieldSchema(name="description", dtype=DataType.VARCHAR, max_length=1000),
|
||||
FieldSchema(name="table_name", dtype=DataType.VARCHAR, max_length=100),
|
||||
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=self.embedding_function.dim)
|
||||
]
|
||||
|
||||
schema = CollectionSchema(fields, description="Text2SQL知识库")
|
||||
self.collection = Collection(name=self.collection_name, schema=schema, consistency_level="Strong")
|
||||
|
||||
index_params = {"index_type": "AUTOINDEX", "metric_type": "IP", "params": {}}
|
||||
self.collection.create_index("embedding", index_params)
|
||||
|
||||
return True
|
||||
|
||||
def load_data(self):
|
||||
"""加载知识库数据"""
|
||||
data_dir = os.path.join(os.path.dirname(__file__), "data")
|
||||
|
||||
self.load_sql_examples(data_dir)
|
||||
self.load_table_schemas(data_dir)
|
||||
self.vectorize_and_store()
|
||||
|
||||
self.data_loaded = True
|
||||
|
||||
def load_sql_examples(self, data_dir: str):
|
||||
"""加载SQL示例"""
|
||||
sql_examples_path = os.path.join(data_dir, "qsql_examples.json")
|
||||
|
||||
default_examples = [
|
||||
{"question": "查询所有用户信息", "sql": "SELECT * FROM users", "description": "获取用户记录", "database": "sqlite"},
|
||||
{"question": "年龄大于30的用户", "sql": "SELECT * FROM users WHERE age > 30", "description": "年龄筛选", "database": "sqlite"},
|
||||
{"question": "统计用户总数", "sql": "SELECT COUNT(*) as user_count FROM users", "description": "用户计数", "database": "sqlite"},
|
||||
{"question": "查询库存不足的产品", "sql": "SELECT * FROM products WHERE stock < 50", "description": "库存筛选", "database": "sqlite"},
|
||||
{"question": "查询用户订单信息", "sql": "SELECT u.name, p.name, o.quantity FROM orders o JOIN users u ON o.user_id = u.id JOIN products p ON o.product_id = p.id", "description": "订单详情", "database": "sqlite"},
|
||||
{"question": "按城市统计用户", "sql": "SELECT city, COUNT(*) as count FROM users GROUP BY city", "description": "城市分组", "database": "sqlite"}
|
||||
]
|
||||
|
||||
if os.path.exists(sql_examples_path):
|
||||
with open(sql_examples_path, 'r', encoding='utf-8') as f:
|
||||
self.sql_examples = json.load(f)
|
||||
else:
|
||||
self.sql_examples = default_examples
|
||||
os.makedirs(data_dir, exist_ok=True)
|
||||
with open(sql_examples_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.sql_examples, f, ensure_ascii=False, indent=2)
|
||||
|
||||
def load_table_schemas(self, data_dir: str):
|
||||
"""加载表结构信息"""
|
||||
schema_path = os.path.join(data_dir, "table_schemas.json")
|
||||
|
||||
default_schemas = [
|
||||
{
|
||||
"table_name": "users",
|
||||
"description": "用户信息表",
|
||||
"columns": [
|
||||
{"name": "id", "type": "INTEGER", "description": "用户ID"},
|
||||
{"name": "name", "type": "VARCHAR", "description": "用户姓名"},
|
||||
{"name": "age", "type": "INTEGER", "description": "用户年龄"},
|
||||
{"name": "email", "type": "VARCHAR", "description": "邮箱地址"},
|
||||
{"name": "city", "type": "VARCHAR", "description": "所在城市"},
|
||||
{"name": "created_at", "type": "DATETIME", "description": "创建时间"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"table_name": "products",
|
||||
"description": "产品信息表",
|
||||
"columns": [
|
||||
{"name": "id", "type": "INTEGER", "description": "产品ID"},
|
||||
{"name": "product_name", "type": "VARCHAR", "description": "产品名称"},
|
||||
{"name": "category", "type": "VARCHAR", "description": "产品类别"},
|
||||
{"name": "price", "type": "DECIMAL", "description": "产品价格"},
|
||||
{"name": "stock", "type": "INTEGER", "description": "库存数量"},
|
||||
{"name": "description", "type": "TEXT", "description": "产品描述"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"table_name": "orders",
|
||||
"description": "订单信息表",
|
||||
"columns": [
|
||||
{"name": "id", "type": "INTEGER", "description": "订单ID"},
|
||||
{"name": "user_id", "type": "INTEGER", "description": "用户ID"},
|
||||
{"name": "product_id", "type": "INTEGER", "description": "产品ID"},
|
||||
{"name": "quantity", "type": "INTEGER", "description": "购买数量"},
|
||||
{"name": "total_price", "type": "DECIMAL", "description": "总价格"},
|
||||
{"name": "order_date", "type": "DATETIME", "description": "订单日期"}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
if os.path.exists(schema_path):
|
||||
with open(schema_path, 'r', encoding='utf-8') as f:
|
||||
self.table_schemas = json.load(f)
|
||||
else:
|
||||
self.table_schemas = default_schemas
|
||||
os.makedirs(data_dir, exist_ok=True)
|
||||
with open(schema_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.table_schemas, f, ensure_ascii=False, indent=2)
|
||||
|
||||
def vectorize_and_store(self):
|
||||
"""向量化数据并存储到Milvus"""
|
||||
self.create_collection()
|
||||
|
||||
all_texts = []
|
||||
all_metadata = []
|
||||
|
||||
for example in self.sql_examples:
|
||||
text = f"问题: {example['question']} SQL: {example['sql']} 描述: {example.get('description', '')}"
|
||||
all_texts.append(text)
|
||||
all_metadata.append({
|
||||
"content_type": "sql_example",
|
||||
"question": example['question'],
|
||||
"sql": example['sql'],
|
||||
"description": example.get('description', ''),
|
||||
"table_name": ""
|
||||
})
|
||||
|
||||
for schema in self.table_schemas:
|
||||
columns_desc = ", ".join([f"{col['name']} ({col['type']}): {col.get('description', '')}"
|
||||
for col in schema['columns']])
|
||||
text = f"表 {schema['table_name']}: {schema['description']} 字段: {columns_desc}"
|
||||
all_texts.append(text)
|
||||
all_metadata.append({
|
||||
"content_type": "table_schema",
|
||||
"question": "",
|
||||
"sql": "",
|
||||
"description": schema['description'],
|
||||
"table_name": schema['table_name']
|
||||
})
|
||||
|
||||
embeddings = self.embedding_function.encode_text(all_texts)
|
||||
|
||||
insert_data = []
|
||||
for i, (embedding, metadata) in enumerate(zip(embeddings, all_metadata)):
|
||||
insert_data.append([
|
||||
metadata["content_type"],
|
||||
metadata["question"],
|
||||
metadata["sql"],
|
||||
metadata["description"],
|
||||
metadata["table_name"],
|
||||
embedding.tolist()
|
||||
])
|
||||
|
||||
self.collection.insert(insert_data)
|
||||
self.collection.flush()
|
||||
self.collection.load()
|
||||
|
||||
def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
|
||||
"""搜索相关的知识库信息"""
|
||||
if not self.data_loaded:
|
||||
self.load_data()
|
||||
|
||||
query_embedding = self.embedding_function.encode_text([query])[0]
|
||||
|
||||
search_params = {"metric_type": "IP", "params": {}}
|
||||
results = self.collection.search(
|
||||
[query_embedding.tolist()],
|
||||
anns_field="embedding",
|
||||
param=search_params,
|
||||
limit=top_k,
|
||||
output_fields=["content_type", "question", "sql", "description", "table_name"]
|
||||
)[0]
|
||||
|
||||
formatted_results = []
|
||||
for hit in results:
|
||||
result = {
|
||||
"score": float(hit.distance),
|
||||
"content_type": hit.entity.get("content_type"),
|
||||
"question": hit.entity.get("question"),
|
||||
"sql": hit.entity.get("sql"),
|
||||
"description": hit.entity.get("description"),
|
||||
"table_name": hit.entity.get("table_name")
|
||||
}
|
||||
formatted_results.append(result)
|
||||
|
||||
return formatted_results
|
||||
|
||||
def _fallback_search(self, query: str, top_k: int) -> List[Dict[str, Any]]:
|
||||
"""降级搜索方法（简单文本匹配）"""
|
||||
results = []
|
||||
query_lower = query.lower()
|
||||
|
||||
for example in self.sql_examples:
|
||||
question_lower = example['question'].lower()
|
||||
sql_lower = example['sql'].lower()
|
||||
|
||||
score = 0
|
||||
for word in query_lower.split():
|
||||
if word in question_lower:
|
||||
score += 2
|
||||
if word in sql_lower:
|
||||
score += 1
|
||||
|
||||
if score > 0:
|
||||
results.append({
|
||||
"score": score,
|
||||
"content_type": "sql_example",
|
||||
"question": example['question'],
|
||||
"sql": example['sql'],
|
||||
"description": example.get('description', ''),
|
||||
"table_name": ""
|
||||
})
|
||||
|
||||
results.sort(key=lambda x: x['score'], reverse=True)
|
||||
return results[:top_k]
|
||||
|
||||
def add_sql_example(self, question: str, sql: str, description: str = ""):
|
||||
"""添加新的SQL示例"""
|
||||
new_example = {
|
||||
"question": question,
|
||||
"sql": sql,
|
||||
"description": description,
|
||||
"database": "sqlite"
|
||||
}
|
||||
self.sql_examples.append(new_example)
|
||||
|
||||
data_dir = os.path.join(os.path.dirname(__file__), "data")
|
||||
sql_examples_path = os.path.join(data_dir, "qsql_examples.json")
|
||||
|
||||
with open(sql_examples_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.sql_examples, f, ensure_ascii=False, indent=2)
|
||||
|
||||
if self.collection and self.data_loaded:
|
||||
text = f"问题: {question} SQL: {sql} 描述: {description}"
|
||||
embedding = self.embedding_function.encode_text([text])[0]
|
||||
|
||||
insert_data = [[
|
||||
"sql_example",
|
||||
question,
|
||||
sql,
|
||||
description,
|
||||
"",
|
||||
embedding.tolist()
|
||||
]]
|
||||
|
||||
self.collection.insert(insert_data)
|
||||
self.collection.flush()
|
||||
|
||||
def cleanup(self):
|
||||
"""清理资源"""
|
||||
if self.collection:
|
||||
self.collection.release()
|
||||
|
||||
if self.milvus_client and self.milvus_client.has_collection(self.collection_name):
|
||||
self.milvus_client.drop_collection(self.collection_name)
|
||||
|
||||
|
||||
def demo():
|
||||
"""简单演示"""
|
||||
# 模型测试
|
||||
embedding_function = BGESmallEmbeddingFunction()
|
||||
test_texts = ["查询用户", "统计数据"]
|
||||
embeddings = embedding_function.encode_text(test_texts)
|
||||
print(f"向量维度: {embeddings.shape}")
|
||||
|
||||
# 数据库查询演示
|
||||
db_path = "demo.db"
|
||||
|
||||
if os.path.exists(db_path):
|
||||
os.remove(db_path)
|
||||
|
||||
conn = sqlite3.connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT, age INTEGER, city TEXT)")
|
||||
|
||||
users_data = [(1, '张三', 25, '北京'), (2, '李四', 32, '上海'), (3, '王五', 35, '深圳')]
|
||||
cursor.executemany("INSERT INTO users VALUES (?, ?, ?, ?)", users_data)
|
||||
|
||||
conn.commit()
|
||||
|
||||
# 执行查询
|
||||
test_sqls = [
|
||||
("查询所有用户", "SELECT * FROM users"),
|
||||
("年龄大于30的用户", "SELECT * FROM users WHERE age > 30"),
|
||||
("统计用户总数", "SELECT COUNT(*) FROM users")
|
||||
]
|
||||
|
||||
for i, (question, sql) in enumerate(test_sqls, 1):
|
||||
print(f"\n问题 {i}: {question}")
|
||||
print("-" * 40)
|
||||
print(f"SQL: {sql}")
|
||||
|
||||
cursor.execute(sql)
|
||||
rows = cursor.fetchall()
|
||||
|
||||
if rows:
|
||||
print(f"返回 {len(rows)} 行数据")
|
||||
for j, row in enumerate(rows[:2], 1):
|
||||
print(f" {j}. {row}")
|
||||
|
||||
if len(rows) > 2:
|
||||
print(f" ... 还有 {len(rows) - 2} 行")
|
||||
else:
|
||||
print("无数据返回")
|
||||
|
||||
conn.close()
|
||||
os.remove(db_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
demo()
|
||||
@@ -0,0 +1,149 @@
|
||||
import os
|
||||
from langchain_deepseek import ChatDeepSeek
|
||||
from langchain_community.document_loaders import BiliBiliLoader
|
||||
from langchain.chains.query_constructor.base import AttributeInfo
|
||||
from openai import OpenAI
|
||||
from langchain_community.vectorstores import Chroma
|
||||
from langchain_huggingface import HuggingFaceEmbeddings
|
||||
import logging
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
# 1. 初始化视频数据
|
||||
video_urls = [
|
||||
"https://www.bilibili.com/video/BV1Bo4y1A7FU",
|
||||
"https://www.bilibili.com/video/BV1ug4y157xA",
|
||||
"https://www.bilibili.com/video/BV1yh411V7ge",
|
||||
]
|
||||
|
||||
bili = []
|
||||
try:
|
||||
loader = BiliBiliLoader(video_urls=video_urls)
|
||||
docs = loader.load()
|
||||
|
||||
for doc in docs:
|
||||
original = doc.metadata
|
||||
|
||||
# 提取基本元数据字段
|
||||
metadata = {
|
||||
'title': original.get('title', '未知标题'),
|
||||
'author': original.get('owner', {}).get('name', '未知作者'),
|
||||
'source': original.get('bvid', '未知ID'),
|
||||
'view_count': original.get('stat', {}).get('view', 0),
|
||||
'length': original.get('duration', 0),
|
||||
}
|
||||
|
||||
doc.metadata = metadata
|
||||
bili.append(doc)
|
||||
|
||||
except Exception as e:
|
||||
print(f"加载BiliBili视频失败: {str(e)}")
|
||||
|
||||
if not bili:
|
||||
print("没有成功加载任何视频，程序退出")
|
||||
exit()
|
||||
|
||||
# 2. 创建向量存储
|
||||
embed_model = HuggingFaceEmbeddings(model_name="BAAI/bge-small-zh-v1.5")
|
||||
vectorstore = Chroma.from_documents(bili, embed_model)
|
||||
|
||||
# 3. 配置元数据字段信息
|
||||
metadata_field_info = [
|
||||
AttributeInfo(
|
||||
name="title",
|
||||
description="视频标题（字符串）",
|
||||
type="string",
|
||||
),
|
||||
AttributeInfo(
|
||||
name="author",
|
||||
description="视频作者（字符串）",
|
||||
type="string",
|
||||
),
|
||||
AttributeInfo(
|
||||
name="view_count",
|
||||
description="视频观看次数（整数）",
|
||||
type="integer",
|
||||
),
|
||||
AttributeInfo(
|
||||
name="length",
|
||||
description="视频长度（整数）",
|
||||
type="integer"
|
||||
)
|
||||
]
|
||||
|
||||
# 4. 初始化LLM客户端
|
||||
client = OpenAI(
|
||||
base_url="https://api.deepseek.com",
|
||||
api_key=os.getenv("DEEPSEEK_API_KEY")
|
||||
)
|
||||
|
||||
# 5. 获取所有文档用于排序
|
||||
all_documents = vectorstore.similarity_search("", k=len(bili))
|
||||
|
||||
# 6. 执行查询示例
|
||||
queries = [
|
||||
"时间最短的视频",
|
||||
"播放量最高的视频"
|
||||
]
|
||||
|
||||
for query in queries:
|
||||
print(f"\n--- 原始查询: '{query}' ---")
|
||||
|
||||
# 使用大模型将自然语言转换为排序指令
|
||||
prompt = f"""你是一个智能助手，请将用户的问题转换成一个用于排序视频的JSON指令。
|
||||
|
||||
你需要识别用户想要排序的字段和排序方向。
|
||||
- 排序字段必须是 'view_count' (观看次数) 或 'length' (时长) 之一。
|
||||
- 排序方向必须是 'asc' (升序) 或 'desc' (降序) 之一。
|
||||
|
||||
例如:
|
||||
- '时间最短的视频' 或 '哪个视频时间最短' 应转换为 {{"sort_by": "length", "order": "asc"}}
|
||||
- '播放量最高的视频' 或 '哪个视频最火' 应转换为 {{"sort_by": "view_count", "order": "desc"}}
|
||||
|
||||
请根据以下问题生成JSON指令:
|
||||
原始问题: "{query}"
|
||||
|
||||
JSON指令:"""
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="deepseek-chat",
|
||||
messages=[
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0,
|
||||
response_format={"type": "json_object"}
|
||||
)
|
||||
|
||||
try:
|
||||
import json
|
||||
instruction_str = response.choices[0].message.content
|
||||
instruction = json.loads(instruction_str)
|
||||
print(f"--- 生成的排序指令: {instruction} ---")
|
||||
|
||||
sort_by = instruction.get('sort_by')
|
||||
order = instruction.get('order')
|
||||
|
||||
if sort_by in ['length', 'view_count'] and order in ['asc', 'desc']:
|
||||
# 在代码中执行排序
|
||||
reverse_order = (order == 'desc')
|
||||
sorted_docs = sorted(all_documents, key=lambda doc: doc.metadata.get(sort_by, 0), reverse=reverse_order)
|
||||
|
||||
# 获取排序后的第一个结果
|
||||
if sorted_docs:
|
||||
doc = sorted_docs[0]
|
||||
title = doc.metadata.get('title', '未知标题')
|
||||
author = doc.metadata.get('author', '未知作者')
|
||||
view_count = doc.metadata.get('view_count', '未知')
|
||||
length = doc.metadata.get('length', '未知')
|
||||
print(f"标题: {title}")
|
||||
print(f"作者: {author}")
|
||||
print(f"观看次数: {view_count}")
|
||||
print(f"时长: {length}秒")
|
||||
print("="*50)
|
||||
else:
|
||||
print("没有找到任何视频")
|
||||
else:
|
||||
print("生成的指令无效，无法执行排序")
|
||||
|
||||
except (json.JSONDecodeError, KeyError) as e:
|
||||
print(f"解析或执行指令失败: {e}")
|
||||
@@ -0,0 +1,72 @@
|
||||
import os
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.output_parsers import StrOutputParser
|
||||
from langchain_deepseek import ChatDeepSeek
|
||||
from langchain_core.runnables import RunnableBranch
|
||||
|
||||
llm = ChatDeepSeek(
|
||||
model="deepseek-chat",
|
||||
temperature=0,
|
||||
api_key=os.getenv("DEEPSEEK_API_KEY")
|
||||
)
|
||||
|
||||
# 1. 设置不同菜系的处理链
|
||||
sichuan_prompt = ChatPromptTemplate.from_template(
|
||||
"你是一位川菜大厨。请用正宗的川菜做法，回答关于「{question}」的问题。"
|
||||
)
|
||||
sichuan_chain = sichuan_prompt | llm | StrOutputParser()
|
||||
|
||||
cantonese_prompt = ChatPromptTemplate.from_template(
|
||||
"你是一位粤菜大厨。请用经典的粤菜做法，回答关于「{question}」的问题。"
|
||||
)
|
||||
cantonese_chain = cantonese_prompt | llm | StrOutputParser()
|
||||
|
||||
# 定义备用通用链
|
||||
general_prompt = ChatPromptTemplate.from_template(
|
||||
"你是一个美食助手。请回答关于「{question}」的问题。"
|
||||
)
|
||||
general_chain = general_prompt | llm | StrOutputParser()
|
||||
|
||||
|
||||
# 2. 创建路由链
|
||||
classifier_prompt = ChatPromptTemplate.from_template(
|
||||
"""根据用户问题中提到的菜品，将其分类为：['川菜', '粤菜', 或 '其他']。
|
||||
不要解释你的理由，只返回一个单词的分类结果。
|
||||
问题: {question}"""
|
||||
)
|
||||
classifier_chain = classifier_prompt | llm | StrOutputParser()
|
||||
|
||||
# 定义路由分支
|
||||
router_branch = RunnableBranch(
|
||||
(lambda x: "川菜" in x["topic"], sichuan_chain),
|
||||
(lambda x: "粤菜" in x["topic"], cantonese_chain),
|
||||
general_chain # 默认选项
|
||||
)
|
||||
|
||||
# 组合成完整路由链
|
||||
full_router_chain = {"topic": classifier_chain, "question": lambda x: x["question"]} | router_branch
|
||||
print("完整的路由链创建成功。\n")
|
||||
|
||||
|
||||
# 3. 运行演示查询
|
||||
demo_questions = [
|
||||
{"question": "麻婆豆腐怎么做？"}, # 应该路由到川菜
|
||||
{"question": "白切鸡的正宗做法是什么？"}, # 应该路由到粤菜
|
||||
{"question": "番茄炒蛋需要放糖吗？"} # 应该路由到其他
|
||||
]
|
||||
|
||||
for i, item in enumerate(demo_questions, 1):
|
||||
question = item["question"]
|
||||
print(f"\n--- 问题 {i}: {question} ---")
|
||||
|
||||
try:
|
||||
# 获取路由决策
|
||||
topic = classifier_chain.invoke({"question": question})
|
||||
print(f"路由决策: {topic}")
|
||||
|
||||
# 执行完整链
|
||||
result = full_router_chain.invoke(item)
|
||||
print(f"回答: {result}")
|
||||
except Exception as e:
|
||||
print(f"执行错误: {e}")
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
import os
|
||||
from langchain_core.prompts import PromptTemplate
|
||||
from langchain_core.output_parsers import StrOutputParser
|
||||
from langchain_deepseek import ChatDeepSeek
|
||||
from langchain_huggingface import HuggingFaceEmbeddings
|
||||
from langchain_core.runnables import RunnableLambda, RunnablePassthrough, RunnablePassthrough
|
||||
from langchain_community.utils.math import cosine_similarity
|
||||
import numpy as np
|
||||
|
||||
# 1. 定义路由描述
|
||||
sichuan_route_prompt = "你是一位处理川菜的专家。用户的问题是关于麻辣、辛香、重口味的菜肴，例如水煮鱼、麻婆豆腐、鱼香肉丝、宫保鸡丁、花椒、海椒等。"
|
||||
cantonese_route_prompt = "你是一位处理粤菜的专家。用户的问题是关于清淡、鲜美、原汁原味的菜肴，例如白切鸡、老火靓汤、虾饺、云吞面等。"
|
||||
|
||||
route_prompts = [sichuan_route_prompt, cantonese_route_prompt]
|
||||
route_names = ["川菜", "粤菜"]
|
||||
|
||||
# 初始化嵌入模型，并对路由描述进行向量化
|
||||
embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-zh-v1.5")
|
||||
route_prompt_embeddings = embeddings.embed_documents(route_prompts)
|
||||
print(f"已定义 {len(route_names)} 个路由: {', '.join(route_names)}")
|
||||
|
||||
# 2. 定义不同路由的目标链
|
||||
llm = ChatDeepSeek(
|
||||
model="deepseek-chat",
|
||||
temperature=0,
|
||||
api_key=os.getenv("DEEPSEEK_API_KEY")
|
||||
)
|
||||
|
||||
# 定义川菜和粤菜处理链
|
||||
sichuan_chain = (
|
||||
PromptTemplate.from_template("你是一位川菜大厨。请用正宗的川菜做法，回答关于「{query}」的问题。")
|
||||
| llm
|
||||
| StrOutputParser()
|
||||
)
|
||||
cantonese_chain = (
|
||||
PromptTemplate.from_template("你是一位粤菜大厨。请用经典的粤菜做法，回答关于「{query}」的问题。")
|
||||
| llm
|
||||
| StrOutputParser()
|
||||
)
|
||||
|
||||
route_map = { "川菜": sichuan_chain, "粤菜": cantonese_chain }
|
||||
print("川菜和粤菜的处理链创建成功。\n")
|
||||
|
||||
# 3. 创建路由函数
|
||||
def route(info):
|
||||
# 对用户查询进行嵌入
|
||||
query_embedding = embeddings.embed_query(info["query"])
|
||||
|
||||
# 计算与各路由提示的余弦相似度
|
||||
similarity_scores = cosine_similarity([query_embedding], route_prompt_embeddings)[0]
|
||||
|
||||
# 找到最相似的路由
|
||||
chosen_route_index = np.argmax(similarity_scores)
|
||||
chosen_route_name = route_names[chosen_route_index]
|
||||
|
||||
print(f"路由决策: 检测到问题与“{chosen_route_name}”最相似。")
|
||||
|
||||
# 获取对应的处理链
|
||||
chosen_chain = route_map[chosen_route_name]
|
||||
|
||||
# 直接调用选中的链并返回结果
|
||||
return chosen_chain.invoke(info)
|
||||
|
||||
# 创建完整的路由链
|
||||
full_chain = RunnableLambda(route)
|
||||
|
||||
|
||||
# 4. 运行演示查询
|
||||
demo_queries = [
|
||||
"水煮鱼怎么做才嫩？", # 应该路由到川菜
|
||||
"如何做一碗清淡的云吞面？", # 应该路由到粤菜
|
||||
"麻婆豆腐的核心调料是什么？", # 应该路由到川菜
|
||||
]
|
||||
|
||||
for i, query in enumerate(demo_queries, 1):
|
||||
print(f"\n--- 问题 {i}: {query} ---")
|
||||
try:
|
||||
# 传入字典，full_chain 会直接返回最终答案
|
||||
result = full_chain.invoke({"query": query})
|
||||
print(f"回答: {result}")
|
||||
except Exception as e:
|
||||
print(f"执行错误: {e}")
|
||||
|
||||
@@ -0,0 +1,186 @@
|
||||
import os
|
||||
from langchain_community.vectorstores import FAISS
|
||||
from langchain.retrievers import ContextualCompressionRetriever
|
||||
from langchain.retrievers.document_compressors import LLMChainExtractor
|
||||
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
from langchain_community.document_loaders import TextLoader
|
||||
from langchain_deepseek import ChatDeepSeek
|
||||
|
||||
# 导入ColBERT重排器需要的模块
|
||||
from langchain.retrievers.document_compressors.base import BaseDocumentCompressor
|
||||
from langchain.retrievers.document_compressors import DocumentCompressorPipeline
|
||||
from langchain_core.documents import Document
|
||||
from typing import Sequence
|
||||
import torch
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
import torch.nn.functional as F
|
||||
|
||||
class ColBERTReranker(BaseDocumentCompressor):
|
||||
"""ColBERT重排器"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
model_name = "bert-base-uncased"
|
||||
|
||||
# 加载模型和分词器
|
||||
object.__setattr__(self, 'tokenizer', AutoTokenizer.from_pretrained(model_name))
|
||||
object.__setattr__(self, 'model', AutoModel.from_pretrained(model_name))
|
||||
self.model.eval()
|
||||
print(f"ColBERT模型加载完成")
|
||||
|
||||
def encode_text(self, texts):
|
||||
"""ColBERT文本编码"""
|
||||
inputs = self.tokenizer(
|
||||
texts,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=128
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = self.model(**inputs)
|
||||
|
||||
embeddings = outputs.last_hidden_state
|
||||
embeddings = F.normalize(embeddings, p=2, dim=-1)
|
||||
|
||||
return embeddings
|
||||
|
||||
def calculate_colbert_similarity(self, query_emb, doc_embs, query_mask, doc_masks):
|
||||
"""ColBERT相似度计算（MaxSim操作）"""
|
||||
scores = []
|
||||
|
||||
for i, doc_emb in enumerate(doc_embs):
|
||||
doc_mask = doc_masks[i:i+1]
|
||||
|
||||
# 计算相似度矩阵
|
||||
similarity_matrix = torch.matmul(query_emb, doc_emb.unsqueeze(0).transpose(-2, -1))
|
||||
|
||||
# 应用文档mask
|
||||
doc_mask_expanded = doc_mask.unsqueeze(1)
|
||||
similarity_matrix = similarity_matrix.masked_fill(~doc_mask_expanded.bool(), -1e9)
|
||||
|
||||
# MaxSim操作
|
||||
max_sim_per_query_token = similarity_matrix.max(dim=-1)[0]
|
||||
|
||||
# 应用查询mask
|
||||
query_mask_expanded = query_mask.unsqueeze(0)
|
||||
max_sim_per_query_token = max_sim_per_query_token.masked_fill(~query_mask_expanded.bool(), 0)
|
||||
|
||||
# 求和得到最终分数
|
||||
colbert_score = max_sim_per_query_token.sum(dim=-1).item()
|
||||
scores.append(colbert_score)
|
||||
|
||||
return scores
|
||||
|
||||
def compress_documents(
|
||||
self,
|
||||
documents: Sequence[Document],
|
||||
query: str,
|
||||
callbacks=None,
|
||||
) -> Sequence[Document]:
|
||||
"""对文档进行ColBERT重排序"""
|
||||
if len(documents) == 0:
|
||||
return documents
|
||||
|
||||
# 编码查询
|
||||
query_inputs = self.tokenizer(
|
||||
[query],
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=128
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
query_outputs = self.model(**query_inputs)
|
||||
query_embeddings = F.normalize(query_outputs.last_hidden_state, p=2, dim=-1)
|
||||
|
||||
# 编码文档
|
||||
doc_texts = [doc.page_content for doc in documents]
|
||||
doc_inputs = self.tokenizer(
|
||||
doc_texts,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=128
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
doc_outputs = self.model(**doc_inputs)
|
||||
doc_embeddings = F.normalize(doc_outputs.last_hidden_state, p=2, dim=-1)
|
||||
|
||||
# 计算ColBERT相似度
|
||||
scores = self.calculate_colbert_similarity(
|
||||
query_embeddings,
|
||||
doc_embeddings,
|
||||
query_inputs['attention_mask'],
|
||||
doc_inputs['attention_mask']
|
||||
)
|
||||
|
||||
# 排序并返回前5个
|
||||
scored_docs = list(zip(documents, scores))
|
||||
scored_docs.sort(key=lambda x: x[1], reverse=True)
|
||||
reranked_docs = [doc for doc, _ in scored_docs[:5]]
|
||||
|
||||
return reranked_docs
|
||||
|
||||
|
||||
|
||||
# 初始化配置
|
||||
hf_bge_embeddings = HuggingFaceBgeEmbeddings(
|
||||
model_name="BAAI/bge-large-zh-v1.5"
|
||||
)
|
||||
|
||||
llm = ChatDeepSeek(
|
||||
model="deepseek-chat",
|
||||
temperature=0.1,
|
||||
api_key=os.getenv("DEEPSEEK_API_KEY")
|
||||
)
|
||||
|
||||
# 1. 加载和处理文档
|
||||
loader = TextLoader("../../data/C4/txt/ai.txt", encoding="utf-8")
|
||||
documents = loader.load()
|
||||
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
|
||||
docs = text_splitter.split_documents(documents)
|
||||
|
||||
# 2. 创建向量存储和基础检索器
|
||||
vectorstore = FAISS.from_documents(docs, hf_bge_embeddings)
|
||||
base_retriever = vectorstore.as_retriever(search_kwargs={"k": 20})
|
||||
|
||||
# 3. 设置ColBERT重排序器
|
||||
reranker = ColBERTReranker()
|
||||
|
||||
# 4. 设置LLM压缩器
|
||||
compressor = LLMChainExtractor.from_llm(llm)
|
||||
|
||||
# 5. 使用DocumentCompressorPipeline组装压缩管道
|
||||
# 流程: ColBERT重排 -> LLM压缩
|
||||
pipeline_compressor = DocumentCompressorPipeline(
|
||||
transformers=[reranker, compressor]
|
||||
)
|
||||
|
||||
# 6. 创建最终的压缩检索器
|
||||
final_retriever = ContextualCompressionRetriever(
|
||||
base_compressor=pipeline_compressor,
|
||||
base_retriever=base_retriever
|
||||
)
|
||||
|
||||
# 7. 执行查询并展示结果
|
||||
query = "AI还有哪些缺陷需要克服？"
|
||||
print(f"\n{'='*20} 开始执行查询 {'='*20}")
|
||||
print(f"查询: {query}\n")
|
||||
|
||||
# 7.1 基础检索结果
|
||||
print(f"--- (1) 基础检索结果 (Top 20) ---")
|
||||
base_results = base_retriever.get_relevant_documents(query)
|
||||
for i, doc in enumerate(base_results):
|
||||
print(f" [{i+1}] {doc.page_content[:100]}...\n")
|
||||
|
||||
# 7.2 使用管道压缩器的最终结果
|
||||
print(f"\n--- (2) 管道压缩后结果 (ColBERT重排 + LLM压缩) ---")
|
||||
final_results = final_retriever.get_relevant_documents(query)
|
||||
for i, doc in enumerate(final_results):
|
||||
print(f" [{i+1}] {doc.page_content}\n")
|
||||
@@ -0,0 +1,17 @@
|
||||
"""
|
||||
简化的Text2SQL框架
|
||||
基于RAGFlow方案实现的Text2SQL框架
|
||||
"""
|
||||
|
||||
__version__ = "1.0.0"
|
||||
__author__ = "RAG Team"
|
||||
|
||||
from .knowledge_base import SimpleKnowledgeBase
|
||||
from .sql_generator import SimpleSQLGenerator
|
||||
from .text2sql_agent import SimpleText2SQLAgent
|
||||
|
||||
__all__ = [
|
||||
"SimpleKnowledgeBase",
|
||||
"SimpleSQLGenerator",
|
||||
"SimpleText2SQLAgent"
|
||||
]
|
||||
@@ -0,0 +1,57 @@
|
||||
[
|
||||
{
|
||||
"table_name": "users",
|
||||
"table_description": "用户信息表，存储注册用户的基本信息",
|
||||
"columns": [
|
||||
{"name": "id", "description": "用户唯一标识符，主键", "type": "INT"},
|
||||
{"name": "name", "description": "用户姓名，不能为空", "type": "VARCHAR(100)"},
|
||||
{"name": "email", "description": "用户邮箱地址，必须唯一", "type": "VARCHAR(150)"},
|
||||
{"name": "age", "description": "用户年龄", "type": "INT"},
|
||||
{"name": "created_at", "description": "用户注册时间", "type": "TIMESTAMP"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"table_name": "orders",
|
||||
"table_description": "订单表，记录用户的购买订单信息",
|
||||
"columns": [
|
||||
{"name": "id", "description": "订单唯一标识符，主键", "type": "INT"},
|
||||
{"name": "user_id", "description": "下单用户的ID，外键关联users表", "type": "INT"},
|
||||
{"name": "product_name", "description": "购买的产品名称", "type": "VARCHAR(200)"},
|
||||
{"name": "quantity", "description": "购买数量", "type": "INT"},
|
||||
{"name": "price", "description": "订单总价格", "type": "DECIMAL(10,2)"},
|
||||
{"name": "order_date", "description": "下单时间", "type": "TIMESTAMP"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"table_name": "products",
|
||||
"table_description": "产品表，存储商城中所有产品的信息",
|
||||
"columns": [
|
||||
{"name": "id", "description": "产品唯一标识符，主键", "type": "INT"},
|
||||
{"name": "name", "description": "产品名称", "type": "VARCHAR(200)"},
|
||||
{"name": "category", "description": "产品分类", "type": "VARCHAR(100)"},
|
||||
{"name": "price", "description": "产品单价", "type": "DECIMAL(10,2)"},
|
||||
{"name": "stock", "description": "库存数量", "type": "INT"},
|
||||
{"name": "description", "description": "产品详细描述", "type": "TEXT"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"table_name": "categories",
|
||||
"table_description": "产品分类表，定义产品的分类信息",
|
||||
"columns": [
|
||||
{"name": "id", "description": "分类唯一标识符，主键", "type": "INT"},
|
||||
{"name": "name", "description": "分类名称，必须唯一", "type": "VARCHAR(100)"},
|
||||
{"name": "description", "description": "分类描述", "type": "TEXT"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"table_name": "order_items",
|
||||
"table_description": "订单明细表，存储订单中包含的具体商品信息",
|
||||
"columns": [
|
||||
{"name": "id", "description": "订单明细唯一标识符，主键", "type": "INT"},
|
||||
{"name": "order_id", "description": "关联的订单ID，外键", "type": "INT"},
|
||||
{"name": "product_id", "description": "关联的产品ID，外键", "type": "INT"},
|
||||
{"name": "quantity", "description": "该商品在订单中的数量", "type": "INT"},
|
||||
{"name": "unit_price", "description": "该商品的单价（下单时的价格）", "type": "DECIMAL(10,2)"}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,27 @@
|
||||
[
|
||||
{
|
||||
"table_name": "users",
|
||||
"ddl_statement": "CREATE TABLE users (id INT PRIMARY KEY AUTO_INCREMENT, name VARCHAR(100) NOT NULL, email VARCHAR(150) UNIQUE, age INT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)",
|
||||
"description": "用户信息表，存储用户基本信息"
|
||||
},
|
||||
{
|
||||
"table_name": "orders",
|
||||
"ddl_statement": "CREATE TABLE orders (id INT PRIMARY KEY AUTO_INCREMENT, user_id INT, product_name VARCHAR(200), quantity INT, price DECIMAL(10,2), order_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (user_id) REFERENCES users(id))",
|
||||
"description": "订单表，存储用户订单信息"
|
||||
},
|
||||
{
|
||||
"table_name": "products",
|
||||
"ddl_statement": "CREATE TABLE products (id INT PRIMARY KEY AUTO_INCREMENT, name VARCHAR(200) NOT NULL, category VARCHAR(100), price DECIMAL(10,2), stock INT DEFAULT 0, description TEXT)",
|
||||
"description": "产品表，存储产品基本信息"
|
||||
},
|
||||
{
|
||||
"table_name": "categories",
|
||||
"ddl_statement": "CREATE TABLE categories (id INT PRIMARY KEY AUTO_INCREMENT, name VARCHAR(100) NOT NULL UNIQUE, description TEXT)",
|
||||
"description": "产品分类表"
|
||||
},
|
||||
{
|
||||
"table_name": "order_items",
|
||||
"ddl_statement": "CREATE TABLE order_items (id INT PRIMARY KEY AUTO_INCREMENT, order_id INT, product_id INT, quantity INT, unit_price DECIMAL(10,2), FOREIGN KEY (order_id) REFERENCES orders(id), FOREIGN KEY (product_id) REFERENCES products(id))",
|
||||
"description": "订单明细表，存储订单中的具体商品信息"
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,62 @@
|
||||
[
|
||||
{
|
||||
"question": "查询所有用户的姓名和邮箱",
|
||||
"sql": "SELECT name, email FROM users",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查找年龄大于25岁的用户",
|
||||
"sql": "SELECT * FROM users WHERE age > 25",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询每个用户的订单数量",
|
||||
"sql": "SELECT u.name, COUNT(o.id) as order_count FROM users u LEFT JOIN orders o ON u.id = o.user_id GROUP BY u.id, u.name",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查找最近7天的订单",
|
||||
"sql": "SELECT * FROM orders WHERE order_date >= DATE_SUB(NOW(), INTERVAL 7 DAY)",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询总销售额最高的前5个产品",
|
||||
"sql": "SELECT p.name, SUM(oi.quantity * oi.unit_price) as total_sales FROM products p JOIN order_items oi ON p.id = oi.product_id GROUP BY p.id, p.name ORDER BY total_sales DESC LIMIT 5",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询某个用户的所有订单",
|
||||
"sql": "SELECT o.*, u.name as user_name FROM orders o JOIN users u ON o.user_id = u.id WHERE u.name = '张三'",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询价格在100到500之间的产品",
|
||||
"sql": "SELECT * FROM products WHERE price BETWEEN 100 AND 500",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询库存少于10的产品",
|
||||
"sql": "SELECT * FROM products WHERE stock < 10",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询每个分类的产品数量",
|
||||
"sql": "SELECT category, COUNT(*) as product_count FROM products GROUP BY category",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询订单总金额大于1000的订单",
|
||||
"sql": "SELECT o.*, SUM(oi.quantity * oi.unit_price) as total_amount FROM orders o JOIN order_items oi ON o.id = oi.order_id GROUP BY o.id HAVING total_amount > 1000",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询没有下过订单的用户",
|
||||
"sql": "SELECT u.* FROM users u LEFT JOIN orders o ON u.id = o.user_id WHERE o.id IS NULL",
|
||||
"database": "ecommerce"
|
||||
},
|
||||
{
|
||||
"question": "查询平均订单金额",
|
||||
"sql": "SELECT AVG(total_amount) as avg_order_amount FROM (SELECT o.id, SUM(oi.quantity * oi.unit_price) as total_amount FROM orders o JOIN order_items oi ON o.id = oi.order_id GROUP BY o.id) as order_totals",
|
||||
"database": "ecommerce"
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,184 @@
|
||||
import json
|
||||
import os
|
||||
from typing import List, Dict, Any
|
||||
from pymilvus import MilvusClient, FieldSchema, CollectionSchema, DataType
|
||||
from pymilvus.model.hybrid import BGEM3EmbeddingFunction
|
||||
|
||||
|
||||
class SimpleKnowledgeBase:
|
||||
"""知识库"""
|
||||
|
||||
def __init__(self, milvus_uri: str = "http://localhost:19530"):
|
||||
self.milvus_uri = milvus_uri
|
||||
self.client = MilvusClient(uri=milvus_uri)
|
||||
self.embedding_function = BGEM3EmbeddingFunction(use_fp16=False, device="cpu")
|
||||
self.collection_name = "text2sql_kb"
|
||||
self._setup_collection()
|
||||
|
||||
def _setup_collection(self):
|
||||
"""设置集合"""
|
||||
if self.client.has_collection(self.collection_name):
|
||||
self.client.drop_collection(self.collection_name)
|
||||
|
||||
# 定义字段
|
||||
fields = [
|
||||
FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100),
|
||||
FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=4096),
|
||||
FieldSchema(name="type", dtype=DataType.VARCHAR, max_length=32), # ddl, qsql, description
|
||||
FieldSchema(name="dense_vector", dtype=DataType.FLOAT_VECTOR, dim=self.embedding_function.dim["dense"])
|
||||
]
|
||||
|
||||
schema = CollectionSchema(fields, description="Text2SQL知识库")
|
||||
|
||||
# 创建集合
|
||||
self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
schema=schema,
|
||||
consistency_level="Strong"
|
||||
)
|
||||
|
||||
# 创建索引
|
||||
index_params = self.client.prepare_index_params()
|
||||
index_params.add_index(
|
||||
field_name="dense_vector",
|
||||
index_type="AUTOINDEX",
|
||||
metric_type="IP"
|
||||
)
|
||||
|
||||
self.client.create_index(
|
||||
collection_name=self.collection_name,
|
||||
index_params=index_params
|
||||
)
|
||||
|
||||
def load_data(self):
|
||||
"""加载所有知识库数据"""
|
||||
data_dir = os.path.join(os.path.dirname(__file__), "data")
|
||||
|
||||
# 加载DDL数据
|
||||
ddl_path = os.path.join(data_dir, "ddl_examples.json")
|
||||
if os.path.exists(ddl_path):
|
||||
with open(ddl_path, 'r', encoding='utf-8') as f:
|
||||
ddl_data = json.load(f)
|
||||
self._add_ddl_data(ddl_data)
|
||||
|
||||
# 加载Q->SQL数据
|
||||
qsql_path = os.path.join(data_dir, "qsql_examples.json")
|
||||
if os.path.exists(qsql_path):
|
||||
with open(qsql_path, 'r', encoding='utf-8') as f:
|
||||
qsql_data = json.load(f)
|
||||
self._add_qsql_data(qsql_data)
|
||||
|
||||
# 加载描述数据
|
||||
desc_path = os.path.join(data_dir, "db_descriptions.json")
|
||||
if os.path.exists(desc_path):
|
||||
with open(desc_path, 'r', encoding='utf-8') as f:
|
||||
desc_data = json.load(f)
|
||||
self._add_description_data(desc_data)
|
||||
|
||||
# 加载集合到内存
|
||||
self.client.load_collection(collection_name=self.collection_name)
|
||||
print("知识库数据加载完成")
|
||||
|
||||
def _add_ddl_data(self, data: List[Dict]):
|
||||
"""添加DDL数据"""
|
||||
contents = []
|
||||
types = []
|
||||
|
||||
for item in data:
|
||||
content = f"表名: {item.get('table_name', '')}\n"
|
||||
content += f"DDL: {item.get('ddl_statement', '')}\n"
|
||||
content += f"描述: {item.get('description', '')}"
|
||||
|
||||
contents.append(content)
|
||||
types.append("ddl")
|
||||
|
||||
self._insert_data(contents, types)
|
||||
|
||||
def _add_qsql_data(self, data: List[Dict]):
|
||||
"""添加Q->SQL数据"""
|
||||
contents = []
|
||||
types = []
|
||||
|
||||
for item in data:
|
||||
content = f"问题: {item.get('question', '')}\n"
|
||||
content += f"SQL: {item.get('sql', '')}"
|
||||
|
||||
contents.append(content)
|
||||
types.append("qsql")
|
||||
|
||||
self._insert_data(contents, types)
|
||||
|
||||
def _add_description_data(self, data: List[Dict]):
|
||||
"""添加描述数据"""
|
||||
contents = []
|
||||
types = []
|
||||
|
||||
for item in data:
|
||||
content = f"表名: {item.get('table_name', '')}\n"
|
||||
content += f"表描述: {item.get('table_description', '')}\n"
|
||||
|
||||
columns = item.get('columns', [])
|
||||
if columns:
|
||||
content += "字段信息:\n"
|
||||
for col in columns:
|
||||
content += f" - {col.get('name', '')}: {col.get('description', '')} ({col.get('type', '')})\n"
|
||||
|
||||
contents.append(content)
|
||||
types.append("description")
|
||||
|
||||
self._insert_data(contents, types)
|
||||
|
||||
def _insert_data(self, contents: List[str], types: List[str]):
|
||||
"""插入数据"""
|
||||
if not contents:
|
||||
return
|
||||
|
||||
# 生成嵌入
|
||||
embeddings = self.embedding_function(contents)
|
||||
|
||||
# 构建插入数据，每一行是一个字典
|
||||
data_to_insert = []
|
||||
for i in range(len(contents)):
|
||||
data_to_insert.append({
|
||||
"content": contents[i],
|
||||
"type": types[i],
|
||||
"dense_vector": embeddings["dense"][i]
|
||||
})
|
||||
|
||||
# 插入数据
|
||||
result = self.client.insert(
|
||||
collection_name=self.collection_name,
|
||||
data=data_to_insert
|
||||
)
|
||||
|
||||
def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
|
||||
"""搜索相关内容"""
|
||||
self.client.load_collection(collection_name=self.collection_name)
|
||||
|
||||
query_embeddings = self.embedding_function([query])
|
||||
|
||||
search_results = self.client.search(
|
||||
collection_name=self.collection_name,
|
||||
data=query_embeddings["dense"],
|
||||
anns_field="dense_vector",
|
||||
search_params={"metric_type": "IP"},
|
||||
limit=top_k,
|
||||
output_fields=["content", "type"]
|
||||
)
|
||||
|
||||
results = []
|
||||
for hit in search_results[0]:
|
||||
results.append({
|
||||
"content": hit["entity"]["content"],
|
||||
"type": hit["entity"]["type"],
|
||||
"score": hit["distance"]
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
def cleanup(self):
|
||||
"""清理资源"""
|
||||
try:
|
||||
self.client.drop_collection(self.collection_name)
|
||||
except:
|
||||
pass
|
||||
@@ -0,0 +1,113 @@
|
||||
import os
|
||||
from typing import List, Dict, Any
|
||||
from langchain_deepseek import ChatDeepSeek
|
||||
from langchain.schema import HumanMessage, SystemMessage
|
||||
|
||||
|
||||
class SimpleSQLGenerator:
|
||||
"""简化的SQL生成器"""
|
||||
|
||||
def __init__(self, api_key: str = None):
|
||||
self.llm = ChatDeepSeek(
|
||||
model="deepseek-chat",
|
||||
temperature=0,
|
||||
api_key=api_key or os.getenv("DEEPSEEK_API_KEY")
|
||||
)
|
||||
|
||||
def generate_sql(self, user_query: str, knowledge_results: List[Dict[str, Any]]) -> str:
|
||||
"""生成SQL语句"""
|
||||
# 构建上下文
|
||||
context = self._build_context(knowledge_results)
|
||||
|
||||
# 构建提示
|
||||
prompt = f"""你是一个SQL专家。请根据以下信息将用户问题转换为SQL查询语句。
|
||||
|
||||
数据库信息：
|
||||
{context}
|
||||
|
||||
用户问题：{user_query}
|
||||
|
||||
要求：
|
||||
1. 只返回SQL语句，不要包含任何解释
|
||||
2. 确保SQL语法正确
|
||||
3. 使用上下文中提供的表名和字段名
|
||||
4. 如果需要JOIN，请根据表结构进行合理关联
|
||||
|
||||
SQL语句："""
|
||||
|
||||
messages = [HumanMessage(content=prompt)]
|
||||
response = self.llm.invoke(messages)
|
||||
|
||||
# 清理SQL语句
|
||||
sql = response.content.strip()
|
||||
if sql.startswith("```sql"):
|
||||
sql = sql[6:]
|
||||
if sql.startswith("```"):
|
||||
sql = sql[3:]
|
||||
if sql.endswith("```"):
|
||||
sql = sql[:-3]
|
||||
|
||||
return sql.strip()
|
||||
|
||||
def fix_sql(self, original_sql: str, error_message: str, knowledge_results: List[Dict[str, Any]]) -> str:
|
||||
"""修复SQL语句"""
|
||||
context = self._build_context(knowledge_results)
|
||||
|
||||
prompt = f"""请修复以下SQL语句的错误。
|
||||
|
||||
数据库信息：
|
||||
{context}
|
||||
|
||||
原始SQL：
|
||||
{original_sql}
|
||||
|
||||
错误信息：
|
||||
{error_message}
|
||||
|
||||
请返回修复后的SQL语句（只返回SQL，不要解释）："""
|
||||
|
||||
messages = [HumanMessage(content=prompt)]
|
||||
response = self.llm.invoke(messages)
|
||||
|
||||
# 清理SQL语句
|
||||
fixed_sql = response.content.strip()
|
||||
if fixed_sql.startswith("```sql"):
|
||||
fixed_sql = fixed_sql[6:]
|
||||
if fixed_sql.startswith("```"):
|
||||
fixed_sql = fixed_sql[3:]
|
||||
if fixed_sql.endswith("```"):
|
||||
fixed_sql = fixed_sql[:-3]
|
||||
|
||||
return fixed_sql.strip()
|
||||
|
||||
def _build_context(self, knowledge_results: List[Dict[str, Any]]) -> str:
|
||||
"""构建上下文信息"""
|
||||
context = ""
|
||||
|
||||
# 按类型分组
|
||||
ddl_info = []
|
||||
qsql_examples = []
|
||||
descriptions = []
|
||||
|
||||
for result in knowledge_results:
|
||||
if result["type"] == "ddl":
|
||||
ddl_info.append(result["content"])
|
||||
elif result["type"] == "qsql":
|
||||
qsql_examples.append(result["content"])
|
||||
elif result["type"] == "description":
|
||||
descriptions.append(result["content"])
|
||||
|
||||
# 构建上下文
|
||||
if ddl_info:
|
||||
context += "=== 表结构信息 ===\n"
|
||||
context += "\n".join(ddl_info) + "\n\n"
|
||||
|
||||
if descriptions:
|
||||
context += "=== 表和字段描述 ===\n"
|
||||
context += "\n".join(descriptions) + "\n\n"
|
||||
|
||||
if qsql_examples:
|
||||
context += "=== 查询示例 ===\n"
|
||||
context += "\n".join(qsql_examples) + "\n\n"
|
||||
|
||||
return context
|
||||
@@ -0,0 +1,213 @@
|
||||
import sqlite3
|
||||
import os
|
||||
from typing import Dict, Any, List, Tuple
|
||||
from .knowledge_base import SimpleKnowledgeBase
|
||||
from .sql_generator import SimpleSQLGenerator
|
||||
|
||||
|
||||
class SimpleText2SQLAgent:
|
||||
"""Text2SQL代理"""
|
||||
|
||||
def __init__(self, milvus_uri: str = "http://localhost:19530", api_key: str = None):
|
||||
"""初始化代理"""
|
||||
self.knowledge_base = SimpleKnowledgeBase(milvus_uri)
|
||||
self.sql_generator = SimpleSQLGenerator(api_key)
|
||||
self.db_path = None
|
||||
self.connection = None
|
||||
|
||||
# 配置参数
|
||||
self.max_retry_count = 3
|
||||
self.top_k_retrieval = 5
|
||||
self.max_result_rows = 100
|
||||
|
||||
def connect_database(self, db_path: str) -> bool:
|
||||
"""连接SQLite数据库"""
|
||||
try:
|
||||
self.db_path = db_path
|
||||
self.connection = sqlite3.connect(db_path)
|
||||
print(f"成功连接到数据库: {db_path}")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"数据库连接失败: {str(e)}")
|
||||
return False
|
||||
|
||||
def load_knowledge_base(self):
|
||||
"""加载知识库"""
|
||||
self.knowledge_base.load_data()
|
||||
|
||||
def query(self, user_question: str) -> Dict[str, Any]:
|
||||
"""执行Text2SQL查询"""
|
||||
if not self.connection:
|
||||
return {
|
||||
"success": False,
|
||||
"error": "数据库未连接",
|
||||
"sql": None,
|
||||
"results": None
|
||||
}
|
||||
|
||||
print(f"\n=== 处理查询: {user_question} ===")
|
||||
|
||||
# 1. 从知识库检索
|
||||
print("检索知识库...")
|
||||
knowledge_results = self.knowledge_base.search(user_question, self.top_k_retrieval)
|
||||
print(f"检索到 {len(knowledge_results)} 条相关信息")
|
||||
|
||||
# 2. 生成SQL
|
||||
print("生成SQL...")
|
||||
sql = self.sql_generator.generate_sql(user_question, knowledge_results)
|
||||
print(f"生成的SQL: {sql}")
|
||||
|
||||
# 3. 执行SQL（带重试）
|
||||
retry_count = 0
|
||||
while retry_count < self.max_retry_count:
|
||||
print(f"执行SQL (尝试 {retry_count + 1}/{self.max_retry_count})...")
|
||||
|
||||
success, result = self._execute_sql(sql)
|
||||
|
||||
if success:
|
||||
print("SQL执行成功!")
|
||||
return {
|
||||
"success": True,
|
||||
"error": None,
|
||||
"sql": sql,
|
||||
"results": result,
|
||||
"retry_count": retry_count
|
||||
}
|
||||
else:
|
||||
print(f"SQL执行失败: {result}")
|
||||
|
||||
if retry_count < self.max_retry_count - 1:
|
||||
print("尝试修复SQL...")
|
||||
sql = self.sql_generator.fix_sql(sql, result, knowledge_results)
|
||||
print(f"修复后的SQL: {sql}")
|
||||
|
||||
retry_count += 1
|
||||
|
||||
return {
|
||||
"success": False,
|
||||
"error": f"超过最大重试次数 ({self.max_retry_count})",
|
||||
"sql": sql,
|
||||
"results": None,
|
||||
"retry_count": retry_count
|
||||
}
|
||||
|
||||
def _execute_sql(self, sql: str) -> Tuple[bool, Any]:
|
||||
"""执行SQL语句"""
|
||||
try:
|
||||
cursor = self.connection.cursor()
|
||||
|
||||
# 添加LIMIT限制
|
||||
if sql.strip().upper().startswith('SELECT') and 'LIMIT' not in sql.upper():
|
||||
sql = f"{sql.rstrip(';')} LIMIT {self.max_result_rows}"
|
||||
|
||||
cursor.execute(sql)
|
||||
|
||||
if sql.strip().upper().startswith('SELECT'):
|
||||
# 查询语句
|
||||
columns = [desc[0] for desc in cursor.description]
|
||||
rows = cursor.fetchall()
|
||||
|
||||
results = []
|
||||
for row in rows:
|
||||
result_row = {}
|
||||
for i, value in enumerate(row):
|
||||
result_row[columns[i]] = value
|
||||
results.append(result_row)
|
||||
|
||||
cursor.close()
|
||||
return True, {
|
||||
"columns": columns,
|
||||
"rows": results,
|
||||
"count": len(results)
|
||||
}
|
||||
else:
|
||||
# 非查询语句
|
||||
self.connection.commit()
|
||||
cursor.close()
|
||||
return True, "SQL执行成功"
|
||||
|
||||
except Exception as e:
|
||||
return False, str(e)
|
||||
|
||||
def add_example(self, question: str, sql: str):
|
||||
"""添加新的Q->SQL示例"""
|
||||
# 简化版本：直接保存到文件
|
||||
data_dir = os.path.join(os.path.dirname(__file__), "data")
|
||||
qsql_path = os.path.join(data_dir, "qsql_examples.json")
|
||||
|
||||
try:
|
||||
import json
|
||||
|
||||
# 读取现有数据
|
||||
if os.path.exists(qsql_path):
|
||||
with open(qsql_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
else:
|
||||
data = []
|
||||
|
||||
# 添加新示例
|
||||
data.append({
|
||||
"question": question,
|
||||
"sql": sql,
|
||||
"database": "sqlite"
|
||||
})
|
||||
|
||||
# 保存
|
||||
with open(qsql_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||||
|
||||
print(f"已添加新示例: {question}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"添加示例失败: {str(e)}")
|
||||
|
||||
def get_table_info(self) -> List[Dict[str, Any]]:
|
||||
"""获取数据库表信息"""
|
||||
if not self.connection:
|
||||
return []
|
||||
|
||||
try:
|
||||
cursor = self.connection.cursor()
|
||||
|
||||
# 获取所有表名
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
|
||||
tables = cursor.fetchall()
|
||||
|
||||
table_info = []
|
||||
for table in tables:
|
||||
table_name = table[0]
|
||||
|
||||
# 获取表结构
|
||||
cursor.execute(f"PRAGMA table_info({table_name})")
|
||||
columns = cursor.fetchall()
|
||||
|
||||
table_info.append({
|
||||
"table_name": table_name,
|
||||
"columns": [
|
||||
{
|
||||
"name": col[1],
|
||||
"type": col[2],
|
||||
"nullable": not col[3],
|
||||
"default": col[4],
|
||||
"primary_key": bool(col[5])
|
||||
}
|
||||
for col in columns
|
||||
]
|
||||
})
|
||||
|
||||
cursor.close()
|
||||
return table_info
|
||||
|
||||
except Exception as e:
|
||||
print(f"获取表信息失败: {str(e)}")
|
||||
return []
|
||||
|
||||
def cleanup(self):
|
||||
"""清理资源"""
|
||||
if self.connection:
|
||||
self.connection.close()
|
||||
self.connection = None
|
||||
print("数据库连接已关闭")
|
||||
|
||||
self.knowledge_base.cleanup()
|
||||
print("知识库已清理")
|
||||
@@ -0,0 +1,193 @@
|
||||
import os
|
||||
from langchain_community.vectorstores import FAISS
|
||||
from langchain.retrievers import ContextualCompressionRetriever
|
||||
from langchain.retrievers.document_compressors import LLMChainExtractor
|
||||
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
from langchain_community.document_loaders import TextLoader
|
||||
from langchain_deepseek import ChatDeepSeek
|
||||
|
||||
# 导入ColBERT重排器需要的模块
|
||||
from langchain.retrievers.document_compressors.base import BaseDocumentCompressor
|
||||
from langchain.retrievers.document_compressors import DocumentCompressorPipeline
|
||||
from langchain_core.documents import Document
|
||||
from typing import Sequence
|
||||
import torch
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
import torch.nn.functional as F
|
||||
|
||||
class ColBERTReranker(BaseDocumentCompressor):
|
||||
"""ColBERT重排器"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
model_name = "bert-base-uncased"
|
||||
|
||||
# 加载模型和分词器
|
||||
object.__setattr__(self, 'tokenizer', AutoTokenizer.from_pretrained(model_name))
|
||||
object.__setattr__(self, 'model', AutoModel.from_pretrained(model_name))
|
||||
self.model.eval()
|
||||
print(f"ColBERT模型加载完成")
|
||||
|
||||
def encode_text(self, texts):
|
||||
"""ColBERT文本编码"""
|
||||
inputs = self.tokenizer(
|
||||
texts,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=128
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = self.model(**inputs)
|
||||
|
||||
embeddings = outputs.last_hidden_state
|
||||
embeddings = F.normalize(embeddings, p=2, dim=-1)
|
||||
|
||||
return embeddings
|
||||
|
||||
def calculate_colbert_similarity(self, query_emb, doc_embs, query_mask, doc_masks):
|
||||
"""ColBERT相似度计算（MaxSim操作）"""
|
||||
scores = []
|
||||
|
||||
for i, doc_emb in enumerate(doc_embs):
|
||||
doc_mask = doc_masks[i:i+1]
|
||||
|
||||
# 计算相似度矩阵
|
||||
similarity_matrix = torch.matmul(query_emb, doc_emb.unsqueeze(0).transpose(-2, -1))
|
||||
|
||||
# 应用文档mask
|
||||
doc_mask_expanded = doc_mask.unsqueeze(1)
|
||||
similarity_matrix = similarity_matrix.masked_fill(~doc_mask_expanded.bool(), -1e9)
|
||||
|
||||
# MaxSim操作
|
||||
max_sim_per_query_token = similarity_matrix.max(dim=-1)[0]
|
||||
|
||||
# 应用查询mask
|
||||
query_mask_expanded = query_mask.unsqueeze(0)
|
||||
max_sim_per_query_token = max_sim_per_query_token.masked_fill(~query_mask_expanded.bool(), 0)
|
||||
|
||||
# 求和得到最终分数
|
||||
colbert_score = max_sim_per_query_token.sum(dim=-1).item()
|
||||
scores.append(colbert_score)
|
||||
|
||||
return scores
|
||||
|
||||
def compress_documents(
|
||||
self,
|
||||
documents: Sequence[Document],
|
||||
query: str,
|
||||
callbacks=None,
|
||||
) -> Sequence[Document]:
|
||||
"""对文档进行ColBERT重排序"""
|
||||
if len(documents) == 0:
|
||||
return documents
|
||||
|
||||
# 编码查询
|
||||
query_inputs = self.tokenizer(
|
||||
[query],
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=128
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
query_outputs = self.model(**query_inputs)
|
||||
query_embeddings = F.normalize(query_outputs.last_hidden_state, p=2, dim=-1)
|
||||
|
||||
# 编码文档
|
||||
doc_texts = [doc.page_content for doc in documents]
|
||||
doc_inputs = self.tokenizer(
|
||||
doc_texts,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=128
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
doc_outputs = self.model(**doc_inputs)
|
||||
doc_embeddings = F.normalize(doc_outputs.last_hidden_state, p=2, dim=-1)
|
||||
|
||||
# 计算ColBERT相似度
|
||||
scores = self.calculate_colbert_similarity(
|
||||
query_embeddings,
|
||||
doc_embeddings,
|
||||
query_inputs['attention_mask'],
|
||||
doc_inputs['attention_mask']
|
||||
)
|
||||
|
||||
# 排序并返回前5个
|
||||
scored_docs = list(zip(documents, scores))
|
||||
scored_docs.sort(key=lambda x: x[1], reverse=True)
|
||||
reranked_docs = [doc for doc, _ in scored_docs[:5]]
|
||||
|
||||
return reranked_docs
|
||||
|
||||
|
||||
|
||||
# 初始化配置
|
||||
hf_bge_embeddings = HuggingFaceBgeEmbeddings(
|
||||
model_name="BAAI/bge-large-zh-v1.5"
|
||||
)
|
||||
|
||||
llm = ChatDeepSeek(
|
||||
model="deepseek-chat",
|
||||
temperature=0.1,
|
||||
api_key=os.getenv("DEEPSEEK_API_KEY")
|
||||
)
|
||||
|
||||
# 1. 加载和处理文档
|
||||
loader = TextLoader("../../data/C4/txt/ai.txt", encoding="utf-8")
|
||||
documents = loader.load()
|
||||
|
||||
# 优化分块策略：减少重叠，使用中文友好的分隔符
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
separators=["\n\n", "\n", "。", "！", "？", "；", "，", " ", ""],
|
||||
chunk_size=300,
|
||||
chunk_overlap=20
|
||||
)
|
||||
|
||||
docs = text_splitter.split_documents(documents)
|
||||
|
||||
# 2. 创建向量存储和基础检索器
|
||||
vectorstore = FAISS.from_documents(docs, hf_bge_embeddings)
|
||||
base_retriever = vectorstore.as_retriever(search_kwargs={"k": 20})
|
||||
|
||||
# 3. 设置ColBERT重排序器
|
||||
reranker = ColBERTReranker()
|
||||
|
||||
# 4. 设置LLM压缩器
|
||||
compressor = LLMChainExtractor.from_llm(llm)
|
||||
|
||||
# 5. 使用DocumentCompressorPipeline组装压缩管道
|
||||
# 流程: ColBERT重排 -> LLM压缩
|
||||
pipeline_compressor = DocumentCompressorPipeline(
|
||||
transformers=[reranker, compressor]
|
||||
)
|
||||
|
||||
# 6. 创建最终的压缩检索器
|
||||
final_retriever = ContextualCompressionRetriever(
|
||||
base_compressor=pipeline_compressor,
|
||||
base_retriever=base_retriever
|
||||
)
|
||||
|
||||
# 7. 执行查询并展示结果
|
||||
query = "AI还有哪些缺陷需要克服？"
|
||||
print(f"\n{'='*20} 开始执行查询 {'='*20}")
|
||||
print(f"查询: {query}\n")
|
||||
|
||||
# 7.1 基础检索结果
|
||||
print(f"--- (1) 基础检索结果 (Top 20) ---")
|
||||
base_results = base_retriever.get_relevant_documents(query)
|
||||
for i, doc in enumerate(base_results):
|
||||
print(f" [{i+1}] {doc.page_content[:100]}...\n")
|
||||
|
||||
# 7.2 使用管道压缩器的最终结果
|
||||
print(f"\n--- (2) 管道压缩后结果 (ColBERT重排 + LLM压缩) ---")
|
||||
final_results = final_retriever.get_relevant_documents(query)
|
||||
for i, doc in enumerate(final_results):
|
||||
print(f" [{i+1}] {doc.page_content}\n")
|
||||
@@ -0,0 +1,50 @@
|
||||
from typing import List
|
||||
import os
|
||||
from langchain_core.prompts import PromptTemplate
|
||||
from pydantic import BaseModel, Field
|
||||
from langchain_core.output_parsers import PydanticOutputParser
|
||||
from langchain_deepseek import ChatDeepSeek
|
||||
|
||||
# 初始化 LLM
|
||||
llm = ChatDeepSeek(
|
||||
model="deepseek-chat",
|
||||
api_key=os.getenv("DEEPSEEK_API_KEY")
|
||||
)
|
||||
|
||||
# 1. 定义数据结构
|
||||
class PersonInfo(BaseModel):
|
||||
name: str = Field(description="人物姓名")
|
||||
age: int = Field(description="人物年龄")
|
||||
skills: List[str] = Field(description="技能列表")
|
||||
|
||||
# 2. 创建解析器
|
||||
parser = PydanticOutputParser(pydantic_object=PersonInfo)
|
||||
|
||||
# 3. 创建提示模板
|
||||
prompt = PromptTemplate(
|
||||
template="请根据以下文本提取信息。\n{format_instructions}\n{text}\n",
|
||||
input_variables=["text"],
|
||||
partial_variables={"format_instructions": parser.get_format_instructions()},
|
||||
)
|
||||
|
||||
# # 打印格式指令
|
||||
# print("\n--- Format Instructions ---")
|
||||
# print(parser.get_format_instructions())
|
||||
# print("--------------------------\n")
|
||||
|
||||
# 4. 创建处理链
|
||||
chain = prompt | llm | parser
|
||||
|
||||
# 5. 定义输入文本并执行调用链
|
||||
text = "张三今年30岁，他擅长Python和Go语言。"
|
||||
result = chain.invoke({"text": text})
|
||||
|
||||
# 6. 打印结果
|
||||
print("\n--- 解析结果 ---")
|
||||
print(f"结果类型: {type(result)}")
|
||||
print(result)
|
||||
print("--------------------\n")
|
||||
|
||||
print(f"姓名: {result.name}")
|
||||
print(f"年龄: {result.age}")
|
||||
print(f"技能: {result.skills}")
|
||||
@@ -0,0 +1,71 @@
|
||||
from openai import OpenAI
|
||||
import os
|
||||
|
||||
# 初始化 OpenAI 客户端
|
||||
client = OpenAI(
|
||||
api_key=os.getenv("DEEPSEEK_API_KEY"),
|
||||
base_url="https://api.deepseek.com",
|
||||
)
|
||||
|
||||
# 定义一个函数，用于发送消息并获取模型的响应
|
||||
def send_messages(messages, tools=None):
|
||||
response = client.chat.completions.create(
|
||||
model="deepseek-chat",
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
tool_choice="auto", # 让模型自主决定是否调用工具
|
||||
)
|
||||
return response.choices[0].message
|
||||
|
||||
# 1. 定义工具（函数）的 Schema
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "获取指定地点的天气信息",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "城市和省份，例如：杭州市, 浙江省",
|
||||
}
|
||||
},
|
||||
"required": ["location"]
|
||||
},
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
# 1. 用户提问，模型决策调用工具
|
||||
messages = [{"role": "user", "content": "杭州今天天气怎么样？"}]
|
||||
print(f"User> {messages[0]['content']}\n")
|
||||
message = send_messages(messages, tools=tools)
|
||||
|
||||
# 2. 执行工具，并将结果返回模型
|
||||
if message.tool_calls:
|
||||
print("--- 模型发起了工具调用 ---")
|
||||
tool_call = message.tool_calls[0]
|
||||
function_info = tool_call.function
|
||||
print(f"工具名称: {function_info.name}")
|
||||
print(f"工具参数: {function_info.arguments}")
|
||||
|
||||
# 将模型的回复（包含工具调用请求）添加到消息历史中
|
||||
messages.append(message)
|
||||
|
||||
# 模拟执行工具
|
||||
tool_output = "24℃，晴朗"
|
||||
print(f"--- 执行工具并返回结果 ---")
|
||||
print(f"工具执行结果: {tool_output}\n")
|
||||
|
||||
# 将工具的执行结果作为一个新的消息添加到历史中
|
||||
messages.append({"role": "tool", "tool_call_id": tool_call.id, "content": tool_output})
|
||||
|
||||
# 3. 第二次调用：将工具结果返回给模型，获取最终回答
|
||||
print("--- 将工具结果返回给模型，获取最终答案 ---")
|
||||
final_message = send_messages(messages, tools=tools)
|
||||
print(f"Model> {final_message.content}")
|
||||
else:
|
||||
# 如果模型没有调用工具，直接打印其回答
|
||||
print(f"Model> {message.content}")
|
||||
@@ -0,0 +1,125 @@
|
||||
import os
|
||||
import asyncio
|
||||
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
|
||||
from llama_index.core.node_parser import SentenceWindowNodeParser, SentenceSplitter
|
||||
from llama_index.llms.deepseek import DeepSeek
|
||||
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
|
||||
from llama_index.core.postprocessor import MetadataReplacementPostProcessor
|
||||
from llama_index.core.evaluation import (
|
||||
FaithfulnessEvaluator,
|
||||
RelevancyEvaluator,
|
||||
BatchEvalRunner,
|
||||
)
|
||||
from llama_index.core.evaluation.eval_utils import get_results_df
|
||||
from llama_index.core.evaluation import DatasetGenerator, QueryResponseDataset
|
||||
|
||||
Settings.llm = DeepSeek(model="deepseek-chat", temperature=0.1, api_key=os.getenv("DEEPSEEK_API_KEY"))
|
||||
Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en")
|
||||
|
||||
async def main():
|
||||
# 1. 加载文档
|
||||
reader = SimpleDirectoryReader(input_files=["../../data/C3/pdf/IPCC_AR6_WGII_Chapter03.pdf"])
|
||||
documents = reader.load_data()
|
||||
|
||||
# 1.1 加载或生成响应评估数据集
|
||||
if os.path.exists("./c6_response_eval_dataset.json"):
|
||||
print("加载响应评估数据集...")
|
||||
response_eval_dataset = QueryResponseDataset.from_json("./c6_response_eval_dataset.json")
|
||||
else:
|
||||
print("生成响应评估数据集...")
|
||||
dataset_generator = DatasetGenerator.from_documents(documents[:10]) # 减少文档数量
|
||||
response_eval_dataset = await dataset_generator.agenerate_dataset_from_nodes(num=15) # 减少问题数量
|
||||
response_eval_dataset.save_json("./c6_response_eval_dataset.json")
|
||||
|
||||
|
||||
|
||||
# 2. 构建两种不同的RAG查询引擎和检索器进行对比
|
||||
# 2.1 句子窗口检索
|
||||
sentence_parser = SentenceWindowNodeParser.from_defaults(
|
||||
window_size=3,
|
||||
window_metadata_key="window",
|
||||
original_text_metadata_key="original_text",
|
||||
)
|
||||
sentence_nodes = sentence_parser.get_nodes_from_documents(documents)
|
||||
sentence_index = VectorStoreIndex(sentence_nodes)
|
||||
|
||||
sentence_query_engine = sentence_index.as_query_engine(
|
||||
similarity_top_k=2,
|
||||
node_postprocessors=[
|
||||
MetadataReplacementPostProcessor(target_metadata_key="window")
|
||||
],
|
||||
)
|
||||
sentence_retriever = sentence_index.as_retriever(similarity_top_k=2)
|
||||
|
||||
# 2.2 常规分块检索（基准）
|
||||
base_parser = SentenceSplitter(chunk_size=512)
|
||||
base_nodes = base_parser.get_nodes_from_documents(documents)
|
||||
base_index = VectorStoreIndex(base_nodes)
|
||||
|
||||
base_query_engine = base_index.as_query_engine(similarity_top_k=2)
|
||||
base_retriever = base_index.as_retriever(similarity_top_k=2)
|
||||
|
||||
# 3. 初始化响应评估器
|
||||
faithfulness_evaluator = FaithfulnessEvaluator(llm=Settings.llm)
|
||||
relevancy_evaluator = RelevancyEvaluator(llm=Settings.llm)
|
||||
|
||||
# 4. 执行响应评估对比
|
||||
print("开始执行响应评估对比...")
|
||||
evaluators = {"faithfulness": faithfulness_evaluator, "relevancy": relevancy_evaluator}
|
||||
queries = response_eval_dataset.queries
|
||||
|
||||
# 句子窗口检索响应评估
|
||||
print("\n=== 评估句子窗口检索 ===")
|
||||
sentence_runner = BatchEvalRunner(evaluators, workers=2, show_progress=True)
|
||||
sentence_response_results = await sentence_runner.aevaluate_queries(
|
||||
queries=queries, query_engine=sentence_query_engine
|
||||
)
|
||||
|
||||
# 常规分块检索响应评估
|
||||
print("\n=== 评估常规分块检索 ===")
|
||||
base_runner = BatchEvalRunner(evaluators, workers=2, show_progress=True)
|
||||
base_response_results = await base_runner.aevaluate_queries(
|
||||
queries=queries, query_engine=base_query_engine
|
||||
)
|
||||
|
||||
# 5. 分析并打印对比结果
|
||||
print("\n" + "="*60)
|
||||
print("响应评估结果对比")
|
||||
print("="*60)
|
||||
|
||||
def calc_response_score(results, metric):
|
||||
if results and results.get(metric):
|
||||
scores = results[metric]
|
||||
return sum(r.passing for r in scores) / len(scores)
|
||||
return 0
|
||||
|
||||
# 句子窗口检索结果
|
||||
sentence_faith = calc_response_score(sentence_response_results, "faithfulness")
|
||||
sentence_rel = calc_response_score(sentence_response_results, "relevancy")
|
||||
|
||||
# 常规分块检索结果
|
||||
base_faith = calc_response_score(base_response_results, "faithfulness")
|
||||
base_rel = calc_response_score(base_response_results, "relevancy")
|
||||
|
||||
print(f"\n句子窗口检索:")
|
||||
print(f" 忠实度: {sentence_faith:.1%}")
|
||||
print(f" 相关性: {sentence_rel:.1%}")
|
||||
|
||||
print(f"\n常规分块检索:")
|
||||
print(f" 忠实度: {base_faith:.1%}")
|
||||
print(f" 相关性: {base_rel:.1%}")
|
||||
|
||||
|
||||
|
||||
# 简单对比
|
||||
if sentence_faith > base_faith and sentence_rel > base_rel:
|
||||
print(f"\n✅ 句子窗口检索在两个维度上都优于常规分块检索")
|
||||
elif sentence_faith > base_faith or sentence_rel > base_rel:
|
||||
print(f"\n⚖️ 句子窗口检索在某些维度上有优势")
|
||||
else:
|
||||
print(f"\n❌ 句子窗口检索未显示明显优势")
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,36 @@
|
||||
{
|
||||
"queries": {
|
||||
"005f4f96-be64-48ad-ab90-38066f27e553": "Here are 10 diverse questions based solely on the provided context information:",
|
||||
"f522fc46-fd0f-462b-9307-d1628b4da112": "**Citation Question**: How should the chapter \"Oceans and Coastal Ecosystems and Their Services\" be cited according to the IPCC AR6 WGII report?",
|
||||
"64581ee3-eca7-4e50-84a0-dfde6f69abda": "**Authorship Question**: Who are the Coordinating Lead Authors of Chapter 3 in the IPCC AR6 WGII report?",
|
||||
"73298e8c-51a6-4424-8d5c-aa6c1ae5eb81": "**Geographical Representation Question**: Which countries are represented by the Lead Authors of this chapter?",
|
||||
"47a221d0-a95e-49f0-a0f8-13bcb2094c13": "**Page Range Question**: What is the page range for Chapter 3 in the IPCC AR6 WGII report?",
|
||||
"64ff4db4-cd98-4741-a35c-e5d3ac8b7cd2": "**DOI Question**: What is the DOI (Digital Object Identifier) for Chapter 3 of the IPCC AR6 WGII report?",
|
||||
"df2772ec-fbb1-4bde-9fde-87ecbbe242ed": "**Contributing Authors Question**: Name three Contributing Authors from different countries listed in the chapter.",
|
||||
"6e4ed553-7dc6-4db0-8c97-f5dfe8540ed3": "**Review Editors Question**: Who are the Review Editors for this chapter, and which countries do they represent?",
|
||||
"30d8e598-4a92-41e8-bafa-f62f503e2417": "**Chapter Scientist Question**: Who is credited as the Chapter Scientist for this chapter, and what is their affiliation?",
|
||||
"dc190ac2-6d6a-496c-a742-4c799cacf9b9": "**Publisher Question**: Which publishers are responsible for the IPCC AR6 WGII report, and where are they based?",
|
||||
"90275ad3-39f0-44b4-ae1a-da47454f4fe8": "**File Metadata Question**: What is the file size and creation date of the PDF document containing Chapter 3?",
|
||||
"2d2edcb3-3f47-4a73-94ff-459290b231c4": "These questions test different aspects of the provided context, including citation format, authorship details, metadata, and structural elements of the report.",
|
||||
"43f7eea1-2b43-4a99-ab50-fba2dd3471db": "Here are 10 diverse quiz/examination questions based strictly on the provided context:",
|
||||
"ea03ffc2-674d-413a-9954-e708c4b0527b": "**Multiple Choice**: According to Section 3.2.2, which of the following is NOT listed as a physical change in the global ocean due to climate change?",
|
||||
"242acbb0-4f32-4732-8cc0-6344f4e0f17c": "a) Sea level rise"
|
||||
},
|
||||
"responses": {
|
||||
"005f4f96-be64-48ad-ab90-38066f27e553": "Here are 10 diverse questions based solely on the provided context information:\n\n1. **Authorship**: Who are the Coordinating Lead Authors of Chapter 3 in the IPCC AR6 WGII report? \n - *Answer*: Sarah R. Cooley (USA) and David S. Schoeman (Australia). \n\n2. **Chapter Title**: What is the full title of Chapter 3 in the IPCC AR6 WGII report? \n - *Answer*: \"Oceans and Coastal Ecosystems and Their Services.\" \n\n3. **Page Range**: On which pages does Chapter 3 appear in the report? \n - *Answer*: Pages 379\u2013550. \n\n4. **DOI**: What is the Digital Object Identifier (DOI) for Chapter 3? \n - *Answer*: doi:10.1017/9781009325844.005. \n\n5. **Contributing Authors**: Name one Contributing Author from Nigeria. \n - *Answer*: Ibidun Adelekan. \n\n6. **Institutional Affiliation**: Which country/institution is Shin-Ichi Ito affiliated with as a Lead Author? \n - *Answer*: Japan. \n\n7. **Chapter Scientist**: Who served as the Chapter Scientist for this chapter? \n - *Answer*: Jessica Bolin (Australia). \n\n8. **Review Editors**: Who were the Review Editors for Chapter 3? \n - *Answer*: Karim Hilmi (Morocco) and Lisa Levin (USA). \n\n9. **Publisher**: Which publishers released the IPCC AR6 WGII report? \n - *Answer*: Cambridge University Press, Cambridge, UK, and New York, NY, USA. \n\n10. **File Metadata**: What is the file size of the PDF document containing Chapter 3? \n - *Answer*: 21,752,444 bytes (or ~21.75 MB). \n\nThese questions are derived exclusively from the provided context, ensuring no external knowledge is used.",
|
||||
"f522fc46-fd0f-462b-9307-d1628b4da112": "The chapter \"Oceans and Coastal Ecosystems and Their Services\" should be cited as follows: \n\n**Cooley, S., D. Schoeman, L. Bopp, P. Boyd, S. Donner, D.Y. Ghebrehiwet, S.-I. Ito, W. Kiessling, P. Martinetto, E. Ojea, M.-F. Racault, B. Rost, and M. Skern-Mauritzen, 2022: Oceans and Coastal Ecosystems and Their Services.** In: *Climate Change 2022: Impacts, Adaptation and Vulnerability. Contribution of Working Group II to the Sixth Assessment Report of the Intergovernmental Panel on Climate Change* [H.-O. P\u00f6rtner, D.C. Roberts, M. Tignor, E.S. Poloczanska, K. Mintenbeck, A. Alegr\u00eda, M. Craig, S. Langsdorf, S. L\u00f6schke, V. M\u00f6ller, A. Okem, B. Rama (eds.)]. Cambridge University Press, Cambridge, UK and New York, NY, USA, pp. 379\u2013550, doi:10.1017/9781009325844.005. \n\nThis citation includes the coordinating lead authors, lead authors, chapter title, report title, editors, publisher details, page range, and DOI.",
|
||||
"64581ee3-eca7-4e50-84a0-dfde6f69abda": "The Coordinating Lead Authors of Chapter 3 (\"Oceans and Coastal Ecosystems and Their Services\") in the IPCC AR6 WGII report are **Sarah R. Cooley (USA)** and **David S. Schoeman (Australia)**. \n\nThis information is explicitly stated in the provided context under the \"Coordinating Lead Authors\" section.",
|
||||
"73298e8c-51a6-4424-8d5c-aa6c1ae5eb81": "Based on the context information provided, the Lead Authors of the chapter \"Oceans and Coastal Ecosystems and Their Services\" represent the following countries: \n\n1. **USA** \u2013 Sarah R. Cooley \n2. **Australia** \u2013 David S. Schoeman \n3. **France** \u2013 Laurent Bopp \n4. **Australia/UK** \u2013 Philip Boyd \n5. **Canada** \u2013 Simon Donner \n6. **Japan** \u2013 Shin-Ichi Ito \n7. **Germany** \u2013 Wolfgang Kiessling \n8. **Argentina** \u2013 Paulina Martinetto \n9. **Spain** \u2013 Elena Ojea \n10. **UK/France** \u2013 Marie-Fanny Racault \n11. **Germany** \u2013 Bj\u00f6rn Rost \n12. **Norway** \u2013 Mette Skern-Mauritzen \n13. **South Africa/Eritrea** \u2013 Dawit Yemane Ghebrehiwet \n\nThus, the countries represented by the Lead Authors are: **USA, Australia, France, UK, Canada, Japan, Germany, Argentina, Spain, Norway, South Africa, and Eritrea**. Some authors have dual affiliations (e.g., Australia/UK, UK/France, South Africa/Eritrea).",
|
||||
"47a221d0-a95e-49f0-a0f8-13bcb2094c13": "The page range for Chapter 3 (\"Oceans and Coastal Ecosystems and Their Services\") in the IPCC AR6 WGII report is **379\u2013550**, as indicated in the citation provided in the context information. \n\nAnswer: The page range for Chapter 3 is 379\u2013550.",
|
||||
"64ff4db4-cd98-4741-a35c-e5d3ac8b7cd2": "The DOI for Chapter 3 of the IPCC AR6 WGII report is **10.1017/9781009325844.005**. \n\nThis is explicitly stated in the citation provided in the context information.",
|
||||
"df2772ec-fbb1-4bde-9fde-87ecbbe242ed": "Here are three Contributing Authors from different countries listed in the chapter:\n\n1. **Johann D. Bell** (Australia) \n2. **William W. L. Cheung** (Canada) \n3. **Elisabeth Holland** (Fiji) \n\nThese authors represent three distinct countries: Australia, Canada, and Fiji.",
|
||||
"6e4ed553-7dc6-4db0-8c97-f5dfe8540ed3": "The Review Editors for this chapter are **Karim Hilmi** (Morocco) and **Lisa Levin** (USA). They represent Morocco and the United States, respectively. \n\n(Answer derived directly from the provided context information.)",
|
||||
"30d8e598-4a92-41e8-bafa-f62f503e2417": "The Chapter Scientist credited for this chapter is **Jessica Bolin**, and their affiliation is **Australia**. \n\n(Source: Context information provided, under the \"Chapter Scientist\" section.)",
|
||||
"dc190ac2-6d6a-496c-a742-4c799cacf9b9": "The IPCC AR6 WGII report is published by two publishers: \n\n1. **Cambridge University Press** (based in Cambridge, UK) \n2. **Cambridge University Press** (based in New York, NY, USA) \n\nThis information is derived from the citation provided in the context: \n*\"Cambridge University Press, Cambridge, UK and New York, NY, USA, pp. 379\u2013550, doi:10.1017/9781009325844.005.\"* \n\n(Note: While the citation lists \"Cambridge University Press\" twice with different locations, it does not specify a separate publisher for the U.S. edition, implying the same publisher operates in both locations.)",
|
||||
"90275ad3-39f0-44b4-ae1a-da47454f4fe8": "The file size of the PDF document containing Chapter 3 is **21,752,444 bytes**, and the creation date is **2025-07-15**. \n\n(Source: Context metadata provided.)",
|
||||
"2d2edcb3-3f47-4a73-94ff-459290b231c4": "Here are the answers to different aspects of the provided context: \n\n1. **Citation Format**: \n The chapter should be cited as: \n *Cooley, S., D. Schoeman, L. Bopp, P. Boyd, S. Donner, D.Y. Ghebrehiwet, S.-I. Ito, W. Kiessling, P. Martinetto, E. Ojea, M.-F. Racault, B. Rost, and M. Skern-Mauritzen, 2022: Oceans and Coastal Ecosystems and Their Services. In: Climate Change 2022: Impacts, Adaptation and Vulnerability. Contribution of Working Group II to the Sixth Assessment Report of the Intergovernmental Panel on Climate Change [H.-O. P\u00f6rtner, D.C. Roberts, M. Tignor, E.S. Poloczanska, K. Mintenbeck, A. Alegr\u00eda, M. Craig, S. Langsdorf, S. L\u00f6schke, V. M\u00f6ller, A. Okem, B. Rama (eds.)]. Cambridge University Press, Cambridge, UK and New York, NY, USA, pp. 379\u2013550, doi:10.1017/9781009325844.005.* \n\n2. **Authorship Details**: \n - **Coordinating Lead Authors**: Sarah R. Cooley (USA), David S. Schoeman (Australia) \n - **Lead Authors**: Laurent Bopp (France), Philip Boyd (Australia/UK), Simon Donner (Canada), Shin-Ichi Ito (Japan), Wolfgang Kiessling (Germany), Paulina Martinetto (Argentina), Elena Ojea (Spain), Marie-Fanny Racault (UK/France), Bj\u00f6rn Rost (Germany), Mette Skern-Mauritzen (Norway), Dawit Yemane Ghebrehiwet (South Africa/Eritrea) \n - **Contributing Authors, Review Editors, and Chapter Scientist** are also listed in the context. \n\n3. **Metadata**: \n - **Page Label**: 379 \n - **File Name**: IPCC_AR6_WGII_Chapter03.pdf \n - **File Path**: ..\\..\\data\\C3\\pdf\\IPCC_AR6_WGII_Chapter03.pdf \n - **File Type**: application/pdf \n - **File Size**: 21,752,444 bytes (~21.7 MB) \n - **Creation Date**: 2025-07-15 \n - **Last Modified Date**: 2025-07-15 \n\n4. **Structural Elements**: \n - The chapter is part of **Working Group II (WGII)** of the **Sixth Assessment Report (AR6)** of the IPCC. \n - It covers **Oceans and Coastal Ecosystems and Their Services** (pages 379\u2013550). \n - The **DOI** for the chapter is **10.1017/9781009325844.005**. \n\nLet me know if you need further details!",
|
||||
"43f7eea1-2b43-4a99-ab50-fba2dd3471db": "Here are 10 diverse quiz/examination questions based strictly on the provided context from Chapter 3 of the IPCC AR6 WGII report: \n\n1. **Multiple Choice**: Which section of Chapter 3 discusses the challenges of attributing changes in marine ecosystems specifically to climate change? \n a) 3.1 Point of Departure \n b) FAQ 3.1 \n c) 3.2.1 Introduction \n d) Box 3.1 \n\n2. **True/False**: According to the table of contents, the chapter includes a FAQ about marine heatwaves and their impacts on marine life and human communities. (True/False) \n\n3. **Short Answer**: Name two types of climate-induced drivers discussed in Section 3.2 (Observed Trends and Projections of Climatic Impact-Drivers in the Global Ocean). \n\n4. **Fill in the Blank**: Section 3.3.4 focuses on __________ and evolutionary adaptation in marine organisms in response to climate drivers. \n\n5. **Matching**: Match the following subsections to their primary focus: \n - 3.2.2 \n - 3.2.3 \n - 3.3.3 \n - 3.4.2 \n Options: \n a) Chemical changes in the ocean \n b) Coastal ecosystems and seas \n c) Physical changes in the ocean \n d) Responses to multiple drivers \n\n6. **Short Answer**: What is the title of Box 3.2, and what organisms does it focus on? \n\n7. **Multiple Choice**: Which FAQ addresses the concept of \"tipping points\" in the ocean? \n a) FAQ 3.1 \n b) FAQ 3.2 \n c) FAQ 3.3 \n d) None of the above \n\n8. **True/False**: Section 3.3.5 discusses ecological responses to single drivers only. (True/False) \n\n9. **Short Answer**: In which section would you find a discussion of acclimation and adaptation mechanisms in marine species? \n\n10. **Essay Prompt (Brief)**: Based on the table of contents, explain how Chapter 3 structures its analysis of climate impacts on oceans and coastal ecosystems (e.g., drivers, biological responses, observed/projected impacts). \n\nThese questions test comprehension of the chapter's structure, key themes, and specific details while adhering strictly to the provided context. Let me know if you'd like adjustments or additional question types!",
|
||||
"ea03ffc2-674d-413a-9954-e708c4b0527b": "To answer this question, we need to refer to **Section 3.2.2 (Physical Changes)** in the provided context. However, the context does not explicitly list the specific physical changes discussed in that section. Since the detailed content of Section 3.2.2 is not included in the given context, we cannot definitively identify which option is **NOT** listed as a physical change.\n\nIf this were a real exam question, you would need to refer to the actual text of **Section 3.2.2** in the IPCC AR6 WGII Chapter 3 document to determine the correct answer. Without that specific information, we cannot provide an accurate response. \n\nWould you like me to infer possible physical changes based on general climate science knowledge? (Note: This would go beyond the given context.)",
|
||||
"242acbb0-4f32-4732-8cc0-6344f4e0f17c": "The context information provided does not explicitly mention \"sea level rise\" in the listed sections or headings. However, given that the chapter focuses on \"Oceans and Coastal Ecosystems and Their Services,\" it is likely that sea level rise is discussed in sections such as **3.2.2 Physical Changes** (which may cover ocean warming, ice melt, and associated sea level changes) or **3.4.2 Coastal Ecosystems and Seas** (which may address impacts on coastal areas). \n\nFor a precise answer, you would need to refer to the content of these sections in the actual document (e.g., page 392 for physical changes or page 410 for coastal ecosystems). Would you like me to infer potential connections based on typical IPCC reporting?"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
"""
|
||||
RAG系统配置文件
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, Any
|
||||
|
||||
@dataclass
|
||||
class RAGConfig:
|
||||
"""RAG系统配置类"""
|
||||
|
||||
# 路径配置
|
||||
data_path: str = "../../data/C8/cook"
|
||||
index_save_path: str = "./vector_index"
|
||||
|
||||
# 模型配置
|
||||
embedding_model: str = "BAAI/bge-small-zh-v1.5"
|
||||
llm_model: str = "kimi-k2-0711-preview"
|
||||
|
||||
# 检索配置
|
||||
top_k: int = 3
|
||||
|
||||
# 生成配置
|
||||
temperature: float = 0.1
|
||||
max_tokens: int = 2048
|
||||
|
||||
def __post_init__(self):
|
||||
"""初始化后的处理"""
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, config_dict: Dict[str, Any]) -> 'RAGConfig':
|
||||
"""从字典创建配置对象"""
|
||||
return cls(**config_dict)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""转换为字典"""
|
||||
return {
|
||||
'data_path': self.data_path,
|
||||
'index_save_path': self.index_save_path,
|
||||
'embedding_model': self.embedding_model,
|
||||
'llm_model': self.llm_model,
|
||||
'top_k': self.top_k,
|
||||
'temperature': self.temperature,
|
||||
'max_tokens': self.max_tokens
|
||||
}
|
||||
|
||||
# 默认配置实例
|
||||
DEFAULT_CONFIG = RAGConfig()
|
||||
+369
@@ -0,0 +1,369 @@
|
||||
"""
|
||||
RAG系统主程序
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
# 添加模块路径
|
||||
sys.path.append(str(Path(__file__).parent))
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from config import DEFAULT_CONFIG, RAGConfig
|
||||
from rag_modules import (
|
||||
DataPreparationModule,
|
||||
IndexConstructionModule,
|
||||
RetrievalOptimizationModule,
|
||||
GenerationIntegrationModule
|
||||
)
|
||||
|
||||
# 加载环境变量
|
||||
load_dotenv()
|
||||
|
||||
# 配置日志
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class RecipeRAGSystem:
|
||||
"""食谱RAG系统主类"""
|
||||
|
||||
def __init__(self, config: RAGConfig = None):
|
||||
"""
|
||||
初始化RAG系统
|
||||
|
||||
Args:
|
||||
config: RAG系统配置，默认使用DEFAULT_CONFIG
|
||||
"""
|
||||
self.config = config or DEFAULT_CONFIG
|
||||
self.data_module = None
|
||||
self.index_module = None
|
||||
self.retrieval_module = None
|
||||
self.generation_module = None
|
||||
|
||||
# 检查数据路径
|
||||
if not Path(self.config.data_path).exists():
|
||||
raise FileNotFoundError(f"数据路径不存在: {self.config.data_path}")
|
||||
|
||||
# 检查API密钥
|
||||
if not os.getenv("MOONSHOT_API_KEY"):
|
||||
raise ValueError("请设置 MOONSHOT_API_KEY 环境变量")
|
||||
|
||||
def initialize_system(self):
|
||||
"""初始化所有模块"""
|
||||
print("🚀 正在初始化RAG系统...")
|
||||
|
||||
# 1. 初始化数据准备模块
|
||||
print("初始化数据准备模块...")
|
||||
self.data_module = DataPreparationModule(self.config.data_path)
|
||||
|
||||
# 2. 初始化索引构建模块
|
||||
print("初始化索引构建模块...")
|
||||
self.index_module = IndexConstructionModule(
|
||||
model_name=self.config.embedding_model,
|
||||
index_save_path=self.config.index_save_path
|
||||
)
|
||||
|
||||
# 3. 初始化生成集成模块
|
||||
print("🤖 初始化生成集成模块...")
|
||||
self.generation_module = GenerationIntegrationModule(
|
||||
model_name=self.config.llm_model,
|
||||
temperature=self.config.temperature,
|
||||
max_tokens=self.config.max_tokens
|
||||
)
|
||||
|
||||
print("✅ 系统初始化完成！")
|
||||
|
||||
def build_knowledge_base(self):
|
||||
"""构建知识库"""
|
||||
print("\n正在构建知识库...")
|
||||
|
||||
# 1. 尝试加载已保存的索引
|
||||
vectorstore = self.index_module.load_index()
|
||||
|
||||
if vectorstore is not None:
|
||||
print("✅ 成功加载已保存的向量索引！")
|
||||
# 仍需要加载文档和分块用于检索模块
|
||||
print("加载食谱文档...")
|
||||
self.data_module.load_documents()
|
||||
print("进行文本分块...")
|
||||
chunks = self.data_module.chunk_documents()
|
||||
else:
|
||||
print("未找到已保存的索引，开始构建新索引...")
|
||||
|
||||
# 2. 加载文档
|
||||
print("加载食谱文档...")
|
||||
self.data_module.load_documents()
|
||||
|
||||
# 3. 文本分块
|
||||
print("进行文本分块...")
|
||||
chunks = self.data_module.chunk_documents()
|
||||
|
||||
# 4. 构建向量索引
|
||||
print("构建向量索引...")
|
||||
vectorstore = self.index_module.build_vector_index(chunks)
|
||||
|
||||
# 5. 保存索引
|
||||
print("保存向量索引...")
|
||||
self.index_module.save_index()
|
||||
|
||||
# 6. 初始化检索优化模块
|
||||
print("初始化检索优化...")
|
||||
self.retrieval_module = RetrievalOptimizationModule(vectorstore, chunks)
|
||||
|
||||
# 7. 显示统计信息
|
||||
stats = self.data_module.get_statistics()
|
||||
print(f"\n📊 知识库统计:")
|
||||
print(f" 文档总数: {stats['total_documents']}")
|
||||
print(f" 文本块数: {stats['total_chunks']}")
|
||||
print(f" 菜品分类: {list(stats['categories'].keys())}")
|
||||
print(f" 难度分布: {stats['difficulties']}")
|
||||
|
||||
print("✅ 知识库构建完成！")
|
||||
|
||||
def ask_question(self, question: str, stream: bool = False):
|
||||
"""
|
||||
回答用户问题
|
||||
|
||||
Args:
|
||||
question: 用户问题
|
||||
stream: 是否使用流式输出
|
||||
|
||||
Returns:
|
||||
生成的回答或生成器
|
||||
"""
|
||||
if not all([self.retrieval_module, self.generation_module]):
|
||||
raise ValueError("请先构建知识库")
|
||||
|
||||
print(f"\n❓ 用户问题: {question}")
|
||||
|
||||
# 1. 查询路由
|
||||
route_type = self.generation_module.query_router(question)
|
||||
print(f"🎯 查询类型: {route_type}")
|
||||
|
||||
# 2. 智能查询重写（根据路由类型）
|
||||
if route_type == 'list':
|
||||
# 列表查询保持原查询
|
||||
rewritten_query = question
|
||||
print(f"📝 列表查询保持原样: {question}")
|
||||
else:
|
||||
# 详细查询和一般查询使用智能重写
|
||||
print("🤖 智能分析查询...")
|
||||
rewritten_query = self.generation_module.query_rewrite(question)
|
||||
|
||||
# 3. 检索相关子块（自动应用元数据过滤）
|
||||
print("🔍 检索相关文档...")
|
||||
filters = self._extract_filters_from_query(question)
|
||||
if filters:
|
||||
print(f"应用过滤条件: {filters}")
|
||||
relevant_chunks = self.retrieval_module.metadata_filtered_search(rewritten_query, filters, top_k=self.config.top_k)
|
||||
else:
|
||||
relevant_chunks = self.retrieval_module.hybrid_search(rewritten_query, top_k=self.config.top_k)
|
||||
|
||||
# 显示检索到的子块信息
|
||||
if relevant_chunks:
|
||||
chunk_info = []
|
||||
for chunk in relevant_chunks:
|
||||
dish_name = chunk.metadata.get('dish_name', '未知菜品')
|
||||
# 尝试从内容中提取章节标题
|
||||
content_preview = chunk.page_content[:100].strip()
|
||||
if content_preview.startswith('#'):
|
||||
# 如果是标题开头，提取标题（仅取第一行）
|
||||
title_end = content_preview.find('\n') if '\n' in content_preview else len(content_preview)
|
||||
section_title = content_preview[:title_end].replace('#', '').strip()
|
||||
chunk_info.append(f"{dish_name}({section_title})")
|
||||
else:
|
||||
chunk_info.append(f"{dish_name}(内容片段)")
|
||||
|
||||
print(f"找到 {len(relevant_chunks)} 个相关文档块: {', '.join(chunk_info)}")
|
||||
else:
|
||||
print(f"找到 {len(relevant_chunks)} 个相关文档块")
|
||||
|
||||
# 4. 检查是否找到相关内容
|
||||
if not relevant_chunks:
|
||||
return "抱歉，没有找到相关的食谱信息。请尝试其他菜品名称或关键词。"
|
||||
|
||||
# 5. 根据路由类型选择回答方式
|
||||
if route_type == 'list':
|
||||
# 列表查询：直接返回菜品名称列表
|
||||
print("📋 生成菜品列表...")
|
||||
relevant_docs = self.data_module.get_parent_documents(relevant_chunks)
|
||||
|
||||
# 显示找到的文档名称
|
||||
doc_names = []
|
||||
for doc in relevant_docs:
|
||||
dish_name = doc.metadata.get('dish_name', '未知菜品')
|
||||
doc_names.append(dish_name)
|
||||
|
||||
if doc_names:
|
||||
print(f"找到文档: {', '.join(doc_names)}")
|
||||
|
||||
return self.generation_module.generate_list_answer(question, relevant_docs)
|
||||
else:
|
||||
# 详细查询：获取完整文档并生成详细回答
|
||||
print("获取完整文档...")
|
||||
relevant_docs = self.data_module.get_parent_documents(relevant_chunks)
|
||||
|
||||
# 显示找到的文档名称
|
||||
doc_names = []
|
||||
for doc in relevant_docs:
|
||||
dish_name = doc.metadata.get('dish_name', '未知菜品')
|
||||
doc_names.append(dish_name)
|
||||
|
||||
if doc_names:
|
||||
print(f"找到文档: {', '.join(doc_names)}")
|
||||
else:
|
||||
print(f"对应 {len(relevant_docs)} 个完整文档")
|
||||
|
||||
print("✍️ 生成详细回答...")
|
||||
|
||||
# 根据路由类型自动选择回答模式
|
||||
if route_type == "detail":
|
||||
# 详细查询使用分步指导模式
|
||||
if stream:
|
||||
return self.generation_module.generate_step_by_step_answer_stream(question, relevant_docs)
|
||||
else:
|
||||
return self.generation_module.generate_step_by_step_answer(question, relevant_docs)
|
||||
else:
|
||||
# 一般查询使用基础回答模式
|
||||
if stream:
|
||||
return self.generation_module.generate_basic_answer_stream(question, relevant_docs)
|
||||
else:
|
||||
return self.generation_module.generate_basic_answer(question, relevant_docs)
|
||||
|
||||
def _extract_filters_from_query(self, query: str) -> dict:
|
||||
"""
|
||||
从用户问题中提取元数据过滤条件
|
||||
"""
|
||||
filters = {}
|
||||
# 分类关键词
|
||||
category_keywords = DataPreparationModule.get_supported_categories()
|
||||
for cat in category_keywords:
|
||||
if cat in query:
|
||||
filters['category'] = cat
|
||||
break
|
||||
|
||||
# 难度关键词
|
||||
difficulty_keywords = DataPreparationModule.get_supported_difficulties()
|
||||
for diff in sorted(difficulty_keywords, key=len, reverse=True):
|
||||
if diff in query:
|
||||
filters['difficulty'] = diff
|
||||
break
|
||||
|
||||
return filters
|
||||
|
||||
def search_by_category(self, category: str, query: str = "") -> List[str]:
|
||||
"""
|
||||
按分类搜索菜品
|
||||
|
||||
Args:
|
||||
category: 菜品分类
|
||||
query: 可选的额外查询条件
|
||||
|
||||
Returns:
|
||||
菜品名称列表
|
||||
"""
|
||||
if not self.retrieval_module:
|
||||
raise ValueError("请先构建知识库")
|
||||
|
||||
# 使用元数据过滤搜索
|
||||
search_query = query if query else category
|
||||
filters = {"category": category}
|
||||
|
||||
docs = self.retrieval_module.metadata_filtered_search(search_query, filters, top_k=10)
|
||||
|
||||
# 提取菜品名称
|
||||
dish_names = []
|
||||
for doc in docs:
|
||||
dish_name = doc.metadata.get('dish_name', '未知菜品')
|
||||
if dish_name not in dish_names:
|
||||
dish_names.append(dish_name)
|
||||
|
||||
return dish_names
|
||||
|
||||
def get_ingredients_list(self, dish_name: str) -> str:
|
||||
"""
|
||||
获取指定菜品的食材信息
|
||||
|
||||
Args:
|
||||
dish_name: 菜品名称
|
||||
|
||||
Returns:
|
||||
食材信息
|
||||
"""
|
||||
if not all([self.retrieval_module, self.generation_module]):
|
||||
raise ValueError("请先构建知识库")
|
||||
|
||||
# 搜索相关文档
|
||||
docs = self.retrieval_module.hybrid_search(dish_name, top_k=3)
|
||||
|
||||
# 生成食材信息
|
||||
answer = self.generation_module.generate_basic_answer(f"{dish_name}需要什么食材？", docs)
|
||||
|
||||
return answer
|
||||
|
||||
def run_interactive(self):
|
||||
"""运行交互式问答"""
|
||||
print("=" * 60)
|
||||
print("🍽️ 尝尝咸淡RAG系统 - 交互式问答 🍽️")
|
||||
print("=" * 60)
|
||||
print("💡 解决您的选择困难症，告别'今天吃什么'的世纪难题！")
|
||||
|
||||
# 初始化系统
|
||||
self.initialize_system()
|
||||
|
||||
# 构建知识库
|
||||
self.build_knowledge_base()
|
||||
|
||||
print("\n交互式问答 (输入'退出'结束):")
|
||||
|
||||
while True:
|
||||
try:
|
||||
user_input = input("\n您的问题: ").strip()
|
||||
if user_input.lower() in ['退出', 'quit', 'exit', '']:
|
||||
break
|
||||
|
||||
# 询问是否使用流式输出
|
||||
stream_choice = input("是否使用流式输出? (y/n, 默认y): ").strip().lower()
|
||||
use_stream = stream_choice != 'n'
|
||||
|
||||
print("\n回答:")
|
||||
if use_stream:
|
||||
# 流式输出
|
||||
for chunk in self.ask_question(user_input, stream=True):
|
||||
print(chunk, end="", flush=True)
|
||||
print("\n")
|
||||
else:
|
||||
# 普通输出
|
||||
answer = self.ask_question(user_input, stream=False)
|
||||
print(f"{answer}\n")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"处理问题时出错: {e}")
|
||||
|
||||
print("\n感谢使用尝尝咸淡RAG系统！")
|
||||
|
||||
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
try:
|
||||
# 创建RAG系统
|
||||
rag_system = RecipeRAGSystem()
|
||||
|
||||
# 运行交互式问答
|
||||
rag_system.run_interactive()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"系统运行出错: {e}")
|
||||
print(f"系统错误: {e}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,13 @@
|
||||
from .data_preparation import DataPreparationModule
|
||||
from .index_construction import IndexConstructionModule
|
||||
from .retrieval_optimization import RetrievalOptimizationModule
|
||||
from .generation_integration import GenerationIntegrationModule
|
||||
|
||||
__all__ = [
|
||||
'DataPreparationModule',
|
||||
'IndexConstructionModule',
|
||||
'RetrievalOptimizationModule',
|
||||
'GenerationIntegrationModule'
|
||||
]
|
||||
|
||||
__version__ = "1.0.0"
|
||||
@@ -0,0 +1,363 @@
|
||||
"""
|
||||
数据准备模块
|
||||
"""
|
||||
|
||||
import logging
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from langchain_text_splitters import MarkdownHeaderTextSplitter
|
||||
from langchain_core.documents import Document
|
||||
from pathlib import Path
|
||||
import uuid
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class DataPreparationModule:
|
||||
"""数据准备模块 - 负责数据加载、清洗和预处理"""
|
||||
# 统一维护的分类与难度配置，供外部复用，避免关键词重复定义
|
||||
CATEGORY_MAPPING = {
|
||||
'meat_dish': '荤菜',
|
||||
'vegetable_dish': '素菜',
|
||||
'soup': '汤品',
|
||||
'dessert': '甜品',
|
||||
'breakfast': '早餐',
|
||||
'staple': '主食',
|
||||
'aquatic': '水产',
|
||||
'condiment': '调料',
|
||||
'drink': '饮品'
|
||||
}
|
||||
CATEGORY_LABELS = list(set(CATEGORY_MAPPING.values()))
|
||||
DIFFICULTY_LABELS = ['非常简单', '简单', '中等', '困难', '非常困难']
|
||||
|
||||
def __init__(self, data_path: str):
|
||||
"""
|
||||
初始化数据准备模块
|
||||
|
||||
Args:
|
||||
data_path: 数据文件夹路径
|
||||
"""
|
||||
self.data_path = data_path
|
||||
self.documents: List[Document] = [] # 父文档（完整食谱）
|
||||
self.chunks: List[Document] = [] # 子文档（按标题分割的小块）
|
||||
self.parent_child_map: Dict[str, str] = {} # 子块ID -> 父文档ID的映射
|
||||
|
||||
def load_documents(self) -> List[Document]:
|
||||
"""
|
||||
加载文档数据
|
||||
|
||||
Returns:
|
||||
加载的文档列表
|
||||
"""
|
||||
logger.info(f"正在从 {self.data_path} 加载文档...")
|
||||
|
||||
# 直接读取Markdown文件以保持原始格式
|
||||
documents = []
|
||||
data_path_obj = Path(self.data_path)
|
||||
|
||||
for md_file in data_path_obj.rglob("*.md"):
|
||||
try:
|
||||
# 直接读取文件内容，保持Markdown格式
|
||||
with open(md_file, 'r', encoding='utf-8') as f:
|
||||
content = f.read()
|
||||
|
||||
# 为每个父文档分配确定性的唯一ID（基于数据根目录的相对路径）
|
||||
try:
|
||||
data_root = Path(self.data_path).resolve()
|
||||
relative_path = Path(md_file).resolve().relative_to(data_root).as_posix()
|
||||
except Exception:
|
||||
relative_path = Path(md_file).as_posix()
|
||||
parent_id = hashlib.md5(relative_path.encode("utf-8")).hexdigest()
|
||||
|
||||
# 创建Document对象
|
||||
doc = Document(
|
||||
page_content=content,
|
||||
metadata={
|
||||
"source": str(md_file),
|
||||
"parent_id": parent_id,
|
||||
"doc_type": "parent" # 标记为父文档
|
||||
}
|
||||
)
|
||||
documents.append(doc)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"读取文件 {md_file} 失败: {e}")
|
||||
|
||||
# 增强文档元数据
|
||||
for doc in documents:
|
||||
self._enhance_metadata(doc)
|
||||
|
||||
self.documents = documents
|
||||
logger.info(f"成功加载 {len(documents)} 个文档")
|
||||
return documents
|
||||
|
||||
def _enhance_metadata(self, doc: Document):
|
||||
"""
|
||||
增强文档元数据
|
||||
|
||||
Args:
|
||||
doc: 需要增强元数据的文档
|
||||
"""
|
||||
file_path = Path(doc.metadata.get('source', ''))
|
||||
path_parts = file_path.parts
|
||||
|
||||
# 提取菜品分类
|
||||
doc.metadata['category'] = '其他'
|
||||
for key, value in self.CATEGORY_MAPPING.items():
|
||||
if key in path_parts:
|
||||
doc.metadata['category'] = value
|
||||
break
|
||||
|
||||
# 提取菜品名称
|
||||
doc.metadata['dish_name'] = file_path.stem
|
||||
|
||||
# 分析难度等级
|
||||
content = doc.page_content
|
||||
if '★★★★★' in content:
|
||||
doc.metadata['difficulty'] = '非常困难'
|
||||
elif '★★★★' in content:
|
||||
doc.metadata['difficulty'] = '困难'
|
||||
elif '★★★' in content:
|
||||
doc.metadata['difficulty'] = '中等'
|
||||
elif '★★' in content:
|
||||
doc.metadata['difficulty'] = '简单'
|
||||
elif '★' in content:
|
||||
doc.metadata['difficulty'] = '非常简单'
|
||||
else:
|
||||
doc.metadata['difficulty'] = '未知'
|
||||
|
||||
@classmethod
|
||||
def get_supported_categories(cls) -> List[str]:
|
||||
"""对外提供支持的分类标签列表"""
|
||||
return cls.CATEGORY_LABELS
|
||||
|
||||
@classmethod
|
||||
def get_supported_difficulties(cls) -> List[str]:
|
||||
"""对外提供支持的难度标签列表"""
|
||||
return cls.DIFFICULTY_LABELS
|
||||
|
||||
def chunk_documents(self) -> List[Document]:
|
||||
"""
|
||||
Markdown结构感知分块
|
||||
|
||||
Returns:
|
||||
分块后的文档列表
|
||||
"""
|
||||
logger.info("正在进行Markdown结构感知分块...")
|
||||
|
||||
if not self.documents:
|
||||
raise ValueError("请先加载文档")
|
||||
|
||||
# 使用Markdown标题分割器
|
||||
chunks = self._markdown_header_split()
|
||||
|
||||
# 为每个chunk添加基础元数据
|
||||
for i, chunk in enumerate(chunks):
|
||||
if 'chunk_id' not in chunk.metadata:
|
||||
# 如果没有chunk_id（比如分割失败的情况），则生成一个
|
||||
chunk.metadata['chunk_id'] = str(uuid.uuid4())
|
||||
chunk.metadata['batch_index'] = i # 在当前批次中的索引
|
||||
chunk.metadata['chunk_size'] = len(chunk.page_content)
|
||||
|
||||
self.chunks = chunks
|
||||
logger.info(f"Markdown分块完成，共生成 {len(chunks)} 个chunk")
|
||||
return chunks
|
||||
|
||||
def _markdown_header_split(self) -> List[Document]:
|
||||
"""
|
||||
使用Markdown标题分割器进行结构化分割
|
||||
|
||||
Returns:
|
||||
按标题结构分割的文档列表
|
||||
"""
|
||||
# 定义要分割的标题层级
|
||||
headers_to_split_on = [
|
||||
("#", "主标题"), # 菜品名称
|
||||
("##", "二级标题"), # 必备原料、计算、操作等
|
||||
("###", "三级标题") # 简易版本、复杂版本等
|
||||
]
|
||||
|
||||
# 创建Markdown分割器
|
||||
markdown_splitter = MarkdownHeaderTextSplitter(
|
||||
headers_to_split_on=headers_to_split_on,
|
||||
strip_headers=False # 保留标题，便于理解上下文
|
||||
)
|
||||
|
||||
all_chunks = []
|
||||
|
||||
for doc in self.documents:
|
||||
try:
|
||||
# 检查文档内容是否包含Markdown标题
|
||||
content_preview = doc.page_content[:200]
|
||||
has_headers = any(line.strip().startswith('#') for line in content_preview.split('\n'))
|
||||
|
||||
if not has_headers:
|
||||
logger.warning(f"文档 {doc.metadata.get('dish_name', '未知')} 内容中没有发现Markdown标题")
|
||||
logger.debug(f"内容预览: {content_preview}")
|
||||
|
||||
# 对每个文档进行Markdown分割
|
||||
md_chunks = markdown_splitter.split_text(doc.page_content)
|
||||
|
||||
logger.debug(f"文档 {doc.metadata.get('dish_name', '未知')} 分割成 {len(md_chunks)} 个chunk")
|
||||
|
||||
# 如果没有分割成功，说明文档可能没有标题结构
|
||||
if len(md_chunks) <= 1:
|
||||
logger.warning(f"文档 {doc.metadata.get('dish_name', '未知')} 未能按标题分割，可能缺少标题结构")
|
||||
|
||||
# 为每个子块建立与父文档的关系
|
||||
parent_id = doc.metadata["parent_id"]
|
||||
|
||||
for i, chunk in enumerate(md_chunks):
|
||||
# 为子块分配唯一ID
|
||||
child_id = str(uuid.uuid4())
|
||||
|
||||
# 合并原文档元数据和新的标题元数据
|
||||
chunk.metadata.update(doc.metadata)
|
||||
chunk.metadata.update({
|
||||
"chunk_id": child_id,
|
||||
"parent_id": parent_id,
|
||||
"doc_type": "child", # 标记为子文档
|
||||
"chunk_index": i # 在父文档中的位置
|
||||
})
|
||||
|
||||
# 建立父子映射关系
|
||||
self.parent_child_map[child_id] = parent_id
|
||||
|
||||
all_chunks.extend(md_chunks)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"文档 {doc.metadata.get('source', '未知')} Markdown分割失败: {e}")
|
||||
# 如果Markdown分割失败，将整个文档作为一个chunk
|
||||
all_chunks.append(doc)
|
||||
|
||||
logger.info(f"Markdown结构分割完成，生成 {len(all_chunks)} 个结构化块")
|
||||
return all_chunks
|
||||
|
||||
def filter_documents_by_category(self, category: str) -> List[Document]:
|
||||
"""
|
||||
按分类过滤文档
|
||||
|
||||
Args:
|
||||
category: 菜品分类
|
||||
|
||||
Returns:
|
||||
过滤后的文档列表
|
||||
"""
|
||||
return [doc for doc in self.documents if doc.metadata.get('category') == category]
|
||||
|
||||
def filter_documents_by_difficulty(self, difficulty: str) -> List[Document]:
|
||||
"""
|
||||
按难度过滤文档
|
||||
|
||||
Args:
|
||||
difficulty: 难度等级
|
||||
|
||||
Returns:
|
||||
过滤后的文档列表
|
||||
"""
|
||||
return [doc for doc in self.documents if doc.metadata.get('difficulty') == difficulty]
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""
|
||||
获取数据统计信息
|
||||
|
||||
Returns:
|
||||
统计信息字典
|
||||
"""
|
||||
if not self.documents:
|
||||
return {}
|
||||
|
||||
categories = {}
|
||||
difficulties = {}
|
||||
|
||||
for doc in self.documents:
|
||||
# 统计分类
|
||||
category = doc.metadata.get('category', '未知')
|
||||
categories[category] = categories.get(category, 0) + 1
|
||||
|
||||
# 统计难度
|
||||
difficulty = doc.metadata.get('difficulty', '未知')
|
||||
difficulties[difficulty] = difficulties.get(difficulty, 0) + 1
|
||||
|
||||
return {
|
||||
'total_documents': len(self.documents),
|
||||
'total_chunks': len(self.chunks),
|
||||
'categories': categories,
|
||||
'difficulties': difficulties,
|
||||
'avg_chunk_size': sum(chunk.metadata.get('chunk_size', 0) for chunk in self.chunks) / len(self.chunks) if self.chunks else 0
|
||||
}
|
||||
|
||||
def export_metadata(self, output_path: str):
|
||||
"""
|
||||
导出元数据到JSON文件
|
||||
|
||||
Args:
|
||||
output_path: 输出文件路径
|
||||
"""
|
||||
import json
|
||||
|
||||
metadata_list = []
|
||||
for doc in self.documents:
|
||||
metadata_list.append({
|
||||
'source': doc.metadata.get('source'),
|
||||
'dish_name': doc.metadata.get('dish_name'),
|
||||
'category': doc.metadata.get('category'),
|
||||
'difficulty': doc.metadata.get('difficulty'),
|
||||
'content_length': len(doc.page_content)
|
||||
})
|
||||
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(metadata_list, f, ensure_ascii=False, indent=2)
|
||||
|
||||
logger.info(f"元数据已导出到: {output_path}")
|
||||
|
||||
def get_parent_documents(self, child_chunks: List[Document]) -> List[Document]:
|
||||
"""
|
||||
根据子块获取对应的父文档（智能去重）
|
||||
|
||||
Args:
|
||||
child_chunks: 检索到的子块列表
|
||||
|
||||
Returns:
|
||||
对应的父文档列表（去重，按相关性排序）
|
||||
"""
|
||||
# 统计每个父文档被匹配的次数（相关性指标）
|
||||
parent_relevance = {}
|
||||
parent_docs_map = {}
|
||||
|
||||
# 收集所有相关的父文档ID和相关性分数
|
||||
for chunk in child_chunks:
|
||||
parent_id = chunk.metadata.get("parent_id")
|
||||
if parent_id:
|
||||
# 增加相关性计数
|
||||
parent_relevance[parent_id] = parent_relevance.get(parent_id, 0) + 1
|
||||
|
||||
# 缓存父文档（避免重复查找）
|
||||
if parent_id not in parent_docs_map:
|
||||
for doc in self.documents:
|
||||
if doc.metadata.get("parent_id") == parent_id:
|
||||
parent_docs_map[parent_id] = doc
|
||||
break
|
||||
|
||||
# 按相关性排序（匹配次数多的排在前面）
|
||||
sorted_parent_ids = sorted(parent_relevance.keys(),
|
||||
key=lambda x: parent_relevance[x],
|
||||
reverse=True)
|
||||
|
||||
# 构建去重后的父文档列表
|
||||
parent_docs = []
|
||||
for parent_id in sorted_parent_ids:
|
||||
if parent_id in parent_docs_map:
|
||||
parent_docs.append(parent_docs_map[parent_id])
|
||||
|
||||
# 收集父文档名称和相关性信息用于日志
|
||||
parent_info = []
|
||||
for doc in parent_docs:
|
||||
dish_name = doc.metadata.get('dish_name', '未知菜品')
|
||||
parent_id = doc.metadata.get('parent_id')
|
||||
relevance_count = parent_relevance.get(parent_id, 0)
|
||||
parent_info.append(f"{dish_name}({relevance_count}块)")
|
||||
|
||||
logger.info(f"从 {len(child_chunks)} 个子块中找到 {len(parent_docs)} 个去重父文档: {', '.join(parent_info)}")
|
||||
return parent_docs
|
||||
@@ -0,0 +1,401 @@
|
||||
"""
|
||||
生成集成模块
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from typing import List
|
||||
|
||||
from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
|
||||
from langchain_community.chat_models.moonshot import MoonshotChat
|
||||
from langchain_core.documents import Document
|
||||
from langchain_core.runnables import RunnablePassthrough
|
||||
from langchain_core.output_parsers import StrOutputParser
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class GenerationIntegrationModule:
|
||||
"""生成集成模块 - 负责LLM集成和回答生成"""
|
||||
|
||||
def __init__(self, model_name: str = "kimi-k2-0711-preview", temperature: float = 0.1, max_tokens: int = 2048):
|
||||
"""
|
||||
初始化生成集成模块
|
||||
|
||||
Args:
|
||||
model_name: 模型名称
|
||||
temperature: 生成温度
|
||||
max_tokens: 最大token数
|
||||
"""
|
||||
self.model_name = model_name
|
||||
self.temperature = temperature
|
||||
self.max_tokens = max_tokens
|
||||
self.llm = None
|
||||
self.setup_llm()
|
||||
|
||||
def setup_llm(self):
|
||||
"""初始化大语言模型"""
|
||||
logger.info(f"正在初始化LLM: {self.model_name}")
|
||||
|
||||
api_key = os.getenv("MOONSHOT_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("请设置 MOONSHOT_API_KEY 环境变量")
|
||||
|
||||
self.llm = MoonshotChat(
|
||||
model=self.model_name,
|
||||
temperature=self.temperature,
|
||||
max_tokens=self.max_tokens,
|
||||
moonshot_api_key=api_key
|
||||
)
|
||||
|
||||
logger.info("LLM初始化完成")
|
||||
|
||||
def generate_basic_answer(self, query: str, context_docs: List[Document]) -> str:
|
||||
"""
|
||||
生成基础回答
|
||||
|
||||
Args:
|
||||
query: 用户查询
|
||||
context_docs: 上下文文档列表
|
||||
|
||||
Returns:
|
||||
生成的回答
|
||||
"""
|
||||
context = self._build_context(context_docs)
|
||||
|
||||
prompt = ChatPromptTemplate.from_template("""
|
||||
你是一位专业的烹饪助手。请根据以下食谱信息回答用户的问题。
|
||||
|
||||
用户问题: {question}
|
||||
|
||||
相关食谱信息:
|
||||
{context}
|
||||
|
||||
请提供详细、实用的回答。如果信息不足，请诚实说明。
|
||||
|
||||
回答:""")
|
||||
|
||||
# 使用LCEL构建链
|
||||
chain = (
|
||||
{"question": RunnablePassthrough(), "context": lambda _: context}
|
||||
| prompt
|
||||
| self.llm
|
||||
| StrOutputParser()
|
||||
)
|
||||
|
||||
response = chain.invoke(query)
|
||||
return response
|
||||
|
||||
def generate_step_by_step_answer(self, query: str, context_docs: List[Document]) -> str:
|
||||
"""
|
||||
生成分步骤回答
|
||||
|
||||
Args:
|
||||
query: 用户查询
|
||||
context_docs: 上下文文档列表
|
||||
|
||||
Returns:
|
||||
分步骤的详细回答
|
||||
"""
|
||||
context = self._build_context(context_docs)
|
||||
|
||||
prompt = ChatPromptTemplate.from_template("""
|
||||
你是一位专业的烹饪导师。请根据食谱信息，为用户提供详细的分步骤指导。
|
||||
|
||||
用户问题: {question}
|
||||
|
||||
相关食谱信息:
|
||||
{context}
|
||||
|
||||
请灵活组织回答，建议包含以下部分（可根据实际内容调整）：
|
||||
|
||||
## 🥘 菜品介绍
|
||||
[简要介绍菜品特点和难度]
|
||||
|
||||
## 🛒 所需食材
|
||||
[列出主要食材和用量]
|
||||
|
||||
## 👨🍳 制作步骤
|
||||
[详细的分步骤说明，每步包含具体操作和大概所需时间]
|
||||
|
||||
## 💡 制作技巧
|
||||
[仅在有实用技巧时包含。优先使用原文中的实用技巧，如果原文的"附加内容"与烹饪无关或为空，可以基于制作步骤总结关键要点，或者完全省略此部分]
|
||||
|
||||
注意：
|
||||
- 根据实际内容灵活调整结构
|
||||
- 不要强行填充无关内容或重复制作步骤中的信息
|
||||
- 重点突出实用性和可操作性
|
||||
- 如果没有额外的技巧要分享，可以省略制作技巧部分
|
||||
|
||||
回答:""")
|
||||
|
||||
chain = (
|
||||
{"question": RunnablePassthrough(), "context": lambda _: context}
|
||||
| prompt
|
||||
| self.llm
|
||||
| StrOutputParser()
|
||||
)
|
||||
|
||||
response = chain.invoke(query)
|
||||
return response
|
||||
|
||||
def query_rewrite(self, query: str) -> str:
|
||||
"""
|
||||
智能查询重写 - 让大模型判断是否需要重写查询
|
||||
|
||||
Args:
|
||||
query: 原始查询
|
||||
|
||||
Returns:
|
||||
重写后的查询或原查询
|
||||
"""
|
||||
prompt = PromptTemplate(
|
||||
template="""
|
||||
你是一个智能查询分析助手。请分析用户的查询，判断是否需要重写以提高食谱搜索效果。
|
||||
|
||||
原始查询: {query}
|
||||
|
||||
分析规则：
|
||||
1. **具体明确的查询**（直接返回原查询）：
|
||||
- 包含具体菜品名称：如"宫保鸡丁怎么做"、"红烧肉的制作方法"
|
||||
- 明确的制作询问：如"蛋炒饭需要什么食材"、"糖醋排骨的步骤"
|
||||
- 具体的烹饪技巧：如"如何炒菜不粘锅"、"怎样调制糖醋汁"
|
||||
|
||||
2. **模糊不清的查询**（需要重写）：
|
||||
- 过于宽泛：如"做菜"、"有什么好吃的"、"推荐个菜"
|
||||
- 缺乏具体信息：如"川菜"、"素菜"、"简单的"
|
||||
- 口语化表达：如"想吃点什么"、"有饮品推荐吗"
|
||||
|
||||
重写原则：
|
||||
- 保持原意不变
|
||||
- 增加相关烹饪术语
|
||||
- 优先推荐简单易做的
|
||||
- 保持简洁性
|
||||
|
||||
示例：
|
||||
- "做菜" → "简单易做的家常菜谱"
|
||||
- "有饮品推荐吗" → "简单饮品制作方法"
|
||||
- "推荐个菜" → "简单家常菜推荐"
|
||||
- "川菜" → "经典川菜菜谱"
|
||||
- "宫保鸡丁怎么做" → "宫保鸡丁怎么做"（保持原查询）
|
||||
- "红烧肉需要什么食材" → "红烧肉需要什么食材"（保持原查询）
|
||||
|
||||
请输出最终查询（如果不需要重写就返回原查询）:""",
|
||||
input_variables=["query"]
|
||||
)
|
||||
|
||||
chain = (
|
||||
{"query": RunnablePassthrough()}
|
||||
| prompt
|
||||
| self.llm
|
||||
| StrOutputParser()
|
||||
)
|
||||
|
||||
response = chain.invoke(query).strip()
|
||||
|
||||
# 记录重写结果
|
||||
if response != query:
|
||||
logger.info(f"查询已重写: '{query}' → '{response}'")
|
||||
else:
|
||||
logger.info(f"查询无需重写: '{query}'")
|
||||
|
||||
return response
|
||||
|
||||
|
||||
|
||||
def query_router(self, query: str) -> str:
|
||||
"""
|
||||
查询路由 - 根据查询类型选择不同的处理方式
|
||||
|
||||
Args:
|
||||
query: 用户查询
|
||||
|
||||
Returns:
|
||||
路由类型 ('list', 'detail', 'general')
|
||||
"""
|
||||
prompt = ChatPromptTemplate.from_template("""
|
||||
根据用户的问题，将其分类为以下三种类型之一：
|
||||
|
||||
1. 'list' - 用户想要获取菜品列表或推荐，只需要菜名
|
||||
例如：推荐几个素菜、有什么川菜、给我3个简单的菜
|
||||
|
||||
2. 'detail' - 用户想要具体的制作方法或详细信息
|
||||
例如：宫保鸡丁怎么做、制作步骤、需要什么食材
|
||||
|
||||
3. 'general' - 其他一般性问题
|
||||
例如：什么是川菜、制作技巧、营养价值
|
||||
|
||||
请只返回分类结果：list、detail 或 general
|
||||
|
||||
用户问题: {query}
|
||||
|
||||
分类结果:""")
|
||||
|
||||
chain = (
|
||||
{"query": RunnablePassthrough()}
|
||||
| prompt
|
||||
| self.llm
|
||||
| StrOutputParser()
|
||||
)
|
||||
|
||||
result = chain.invoke(query).strip().lower()
|
||||
|
||||
# 确保返回有效的路由类型
|
||||
if result in ['list', 'detail', 'general']:
|
||||
return result
|
||||
else:
|
||||
return 'general' # 默认类型
|
||||
|
||||
def generate_list_answer(self, query: str, context_docs: List[Document]) -> str:
|
||||
"""
|
||||
生成列表式回答 - 适用于推荐类查询
|
||||
|
||||
Args:
|
||||
query: 用户查询
|
||||
context_docs: 上下文文档列表
|
||||
|
||||
Returns:
|
||||
列表式回答
|
||||
"""
|
||||
if not context_docs:
|
||||
return "抱歉，没有找到相关的菜品信息。"
|
||||
|
||||
# 提取菜品名称
|
||||
dish_names = []
|
||||
for doc in context_docs:
|
||||
dish_name = doc.metadata.get('dish_name', '未知菜品')
|
||||
if dish_name not in dish_names:
|
||||
dish_names.append(dish_name)
|
||||
|
||||
# 构建简洁的列表回答
|
||||
if len(dish_names) == 1:
|
||||
return f"为您推荐：{dish_names[0]}"
|
||||
elif len(dish_names) <= 3:
|
||||
return f"为您推荐以下菜品：\n" + "\n".join([f"{i+1}. {name}" for i, name in enumerate(dish_names)])
|
||||
else:
|
||||
return f"为您推荐以下菜品：\n" + "\n".join([f"{i+1}. {name}" for i, name in enumerate(dish_names[:3])]) + f"\n\n还有其他 {len(dish_names)-3} 道菜品可供选择。"
|
||||
|
||||
def generate_basic_answer_stream(self, query: str, context_docs: List[Document]):
|
||||
"""
|
||||
生成基础回答 - 流式输出
|
||||
|
||||
Args:
|
||||
query: 用户查询
|
||||
context_docs: 上下文文档列表
|
||||
|
||||
Yields:
|
||||
生成的回答片段
|
||||
"""
|
||||
context = self._build_context(context_docs)
|
||||
|
||||
prompt = ChatPromptTemplate.from_template("""
|
||||
你是一位专业的烹饪助手。请根据以下食谱信息回答用户的问题。
|
||||
|
||||
用户问题: {question}
|
||||
|
||||
相关食谱信息:
|
||||
{context}
|
||||
|
||||
请提供详细、实用的回答。如果信息不足，请诚实说明。
|
||||
|
||||
回答:""")
|
||||
|
||||
chain = (
|
||||
{"question": RunnablePassthrough(), "context": lambda _: context}
|
||||
| prompt
|
||||
| self.llm
|
||||
| StrOutputParser()
|
||||
)
|
||||
|
||||
for chunk in chain.stream(query):
|
||||
yield chunk
|
||||
|
||||
def generate_step_by_step_answer_stream(self, query: str, context_docs: List[Document]):
|
||||
"""
|
||||
生成详细步骤回答 - 流式输出
|
||||
|
||||
Args:
|
||||
query: 用户查询
|
||||
context_docs: 上下文文档列表
|
||||
|
||||
Yields:
|
||||
详细步骤回答片段
|
||||
"""
|
||||
context = self._build_context(context_docs)
|
||||
|
||||
prompt = ChatPromptTemplate.from_template("""
|
||||
你是一位专业的烹饪导师。请根据食谱信息，为用户提供详细的分步骤指导。
|
||||
|
||||
用户问题: {question}
|
||||
|
||||
相关食谱信息:
|
||||
{context}
|
||||
|
||||
请灵活组织回答，建议包含以下部分（可根据实际内容调整）：
|
||||
|
||||
## 🥘 菜品介绍
|
||||
[简要介绍菜品特点和难度]
|
||||
|
||||
## 🛒 所需食材
|
||||
[列出主要食材和用量]
|
||||
|
||||
## 👨🍳 制作步骤
|
||||
[详细的分步骤说明，每步包含具体操作和大概所需时间]
|
||||
|
||||
## 💡 制作技巧
|
||||
[仅在有实用技巧时包含。如果原文的"附加内容"与烹饪无关或为空，可以基于制作步骤总结关键要点，或者完全省略此部分]
|
||||
|
||||
注意：
|
||||
- 根据实际内容灵活调整结构
|
||||
- 不要强行填充无关内容
|
||||
- 重点突出实用性和可操作性
|
||||
|
||||
回答:""")
|
||||
|
||||
chain = (
|
||||
{"question": RunnablePassthrough(), "context": lambda _: context}
|
||||
| prompt
|
||||
| self.llm
|
||||
| StrOutputParser()
|
||||
)
|
||||
|
||||
for chunk in chain.stream(query):
|
||||
yield chunk
|
||||
|
||||
def _build_context(self, docs: List[Document], max_length: int = 2000) -> str:
|
||||
"""
|
||||
构建上下文字符串
|
||||
|
||||
Args:
|
||||
docs: 文档列表
|
||||
max_length: 最大长度
|
||||
|
||||
Returns:
|
||||
格式化的上下文字符串
|
||||
"""
|
||||
if not docs:
|
||||
return "暂无相关食谱信息。"
|
||||
|
||||
context_parts = []
|
||||
current_length = 0
|
||||
|
||||
for i, doc in enumerate(docs, 1):
|
||||
# 添加元数据信息
|
||||
metadata_info = f"【食谱 {i}】"
|
||||
if 'dish_name' in doc.metadata:
|
||||
metadata_info += f" {doc.metadata['dish_name']}"
|
||||
if 'category' in doc.metadata:
|
||||
metadata_info += f" | 分类: {doc.metadata['category']}"
|
||||
if 'difficulty' in doc.metadata:
|
||||
metadata_info += f" | 难度: {doc.metadata['difficulty']}"
|
||||
|
||||
# 构建文档文本
|
||||
doc_text = f"{metadata_info}\n{doc.page_content}\n"
|
||||
|
||||
# 检查长度限制
|
||||
if current_length + len(doc_text) > max_length:
|
||||
break
|
||||
|
||||
context_parts.append(doc_text)
|
||||
current_length += len(doc_text)
|
||||
|
||||
return "\n" + "="*50 + "\n".join(context_parts)
|
||||
@@ -0,0 +1,135 @@
|
||||
"""
|
||||
索引构建模块
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import List
|
||||
from pathlib import Path
|
||||
|
||||
from langchain_huggingface import HuggingFaceEmbeddings
|
||||
from langchain_community.vectorstores import FAISS
|
||||
from langchain_core.documents import Document
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class IndexConstructionModule:
|
||||
"""索引构建模块 - 负责向量化和索引构建"""
|
||||
|
||||
def __init__(self, model_name: str = "BAAI/bge-small-zh-v1.5", index_save_path: str = "./vector_index"):
|
||||
"""
|
||||
初始化索引构建模块
|
||||
|
||||
Args:
|
||||
model_name: 嵌入模型名称
|
||||
index_save_path: 索引保存路径
|
||||
"""
|
||||
self.model_name = model_name
|
||||
self.index_save_path = index_save_path
|
||||
self.embeddings = None
|
||||
self.vectorstore = None
|
||||
self.setup_embeddings()
|
||||
|
||||
def setup_embeddings(self):
|
||||
"""初始化嵌入模型"""
|
||||
logger.info(f"正在初始化嵌入模型: {self.model_name}")
|
||||
|
||||
self.embeddings = HuggingFaceEmbeddings(
|
||||
model_name=self.model_name,
|
||||
model_kwargs={'device': 'cpu'},
|
||||
encode_kwargs={'normalize_embeddings': True}
|
||||
)
|
||||
|
||||
logger.info("嵌入模型初始化完成")
|
||||
|
||||
def build_vector_index(self, chunks: List[Document]) -> FAISS:
|
||||
"""
|
||||
构建向量索引
|
||||
|
||||
Args:
|
||||
chunks: 文档块列表
|
||||
|
||||
Returns:
|
||||
FAISS向量存储对象
|
||||
"""
|
||||
logger.info("正在构建FAISS向量索引...")
|
||||
|
||||
if not chunks:
|
||||
raise ValueError("文档块列表不能为空")
|
||||
|
||||
# 构建FAISS向量存储
|
||||
self.vectorstore = FAISS.from_documents(
|
||||
documents=chunks,
|
||||
embedding=self.embeddings
|
||||
)
|
||||
|
||||
logger.info(f"向量索引构建完成，包含 {len(chunks)} 个向量")
|
||||
return self.vectorstore
|
||||
|
||||
def add_documents(self, new_chunks: List[Document]):
|
||||
"""
|
||||
向现有索引添加新文档
|
||||
|
||||
Args:
|
||||
new_chunks: 新的文档块列表
|
||||
"""
|
||||
if not self.vectorstore:
|
||||
raise ValueError("请先构建向量索引")
|
||||
|
||||
logger.info(f"正在添加 {len(new_chunks)} 个新文档到索引...")
|
||||
self.vectorstore.add_documents(new_chunks)
|
||||
logger.info("新文档添加完成")
|
||||
|
||||
def save_index(self):
|
||||
"""
|
||||
保存向量索引到配置的路径
|
||||
"""
|
||||
if not self.vectorstore:
|
||||
raise ValueError("请先构建向量索引")
|
||||
|
||||
# 确保保存目录存在
|
||||
Path(self.index_save_path).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.vectorstore.save_local(self.index_save_path)
|
||||
logger.info(f"向量索引已保存到: {self.index_save_path}")
|
||||
|
||||
def load_index(self):
|
||||
"""
|
||||
从配置的路径加载向量索引
|
||||
|
||||
Returns:
|
||||
加载的向量存储对象，如果加载失败返回None
|
||||
"""
|
||||
if not self.embeddings:
|
||||
self.setup_embeddings()
|
||||
|
||||
if not Path(self.index_save_path).exists():
|
||||
logger.info(f"索引路径不存在: {self.index_save_path}，将构建新索引")
|
||||
return None
|
||||
|
||||
try:
|
||||
self.vectorstore = FAISS.load_local(
|
||||
self.index_save_path,
|
||||
self.embeddings,
|
||||
allow_dangerous_deserialization=True
|
||||
)
|
||||
logger.info(f"向量索引已从 {self.index_save_path} 加载")
|
||||
return self.vectorstore
|
||||
except Exception as e:
|
||||
logger.warning(f"加载向量索引失败: {e}，将构建新索引")
|
||||
return None
|
||||
|
||||
def similarity_search(self, query: str, k: int = 5) -> List[Document]:
|
||||
"""
|
||||
相似度搜索
|
||||
|
||||
Args:
|
||||
query: 查询文本
|
||||
k: 返回结果数量
|
||||
|
||||
Returns:
|
||||
相似文档列表
|
||||
"""
|
||||
if not self.vectorstore:
|
||||
raise ValueError("请先构建或加载向量索引")
|
||||
|
||||
return self.vectorstore.similarity_search(query, k=k)
|
||||
@@ -0,0 +1,162 @@
|
||||
"""
|
||||
检索优化模块
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from langchain_community.vectorstores import FAISS
|
||||
from langchain_community.retrievers import BM25Retriever
|
||||
from langchain_core.documents import Document
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class RetrievalOptimizationModule:
|
||||
"""检索优化模块 - 负责混合检索和过滤"""
|
||||
|
||||
def __init__(self, vectorstore: FAISS, chunks: List[Document]):
|
||||
"""
|
||||
初始化检索优化模块
|
||||
|
||||
Args:
|
||||
vectorstore: FAISS向量存储
|
||||
chunks: 文档块列表
|
||||
"""
|
||||
self.vectorstore = vectorstore
|
||||
self.chunks = chunks
|
||||
self.setup_retrievers()
|
||||
|
||||
def setup_retrievers(self):
|
||||
"""设置向量检索器和BM25检索器"""
|
||||
logger.info("正在设置检索器...")
|
||||
|
||||
# 向量检索器
|
||||
self.vector_retriever = self.vectorstore.as_retriever(
|
||||
search_type="similarity",
|
||||
search_kwargs={"k": 5}
|
||||
)
|
||||
|
||||
# BM25检索器
|
||||
self.bm25_retriever = BM25Retriever.from_documents(
|
||||
self.chunks,
|
||||
k=5
|
||||
)
|
||||
|
||||
|
||||
|
||||
logger.info("检索器设置完成")
|
||||
|
||||
def hybrid_search(self, query: str, top_k: int = 3) -> List[Document]:
|
||||
"""
|
||||
混合检索 - 结合向量检索和BM25检索，使用RRF重排
|
||||
|
||||
Args:
|
||||
query: 查询文本
|
||||
top_k: 返回结果数量
|
||||
|
||||
Returns:
|
||||
检索到的文档列表
|
||||
"""
|
||||
# 分别获取向量检索和BM25检索结果
|
||||
vector_docs = self.vector_retriever.invoke(query)
|
||||
bm25_docs = self.bm25_retriever.invoke(query)
|
||||
|
||||
# 使用RRF重排
|
||||
reranked_docs = self._rrf_rerank(vector_docs, bm25_docs)
|
||||
return reranked_docs[:top_k]
|
||||
|
||||
def metadata_filtered_search(self, query: str, filters: Dict[str, Any], top_k: int = 5) -> List[Document]:
|
||||
"""
|
||||
带元数据过滤的检索
|
||||
|
||||
Args:
|
||||
query: 查询文本
|
||||
filters: 元数据过滤条件
|
||||
top_k: 返回结果数量
|
||||
|
||||
Returns:
|
||||
过滤后的文档列表
|
||||
"""
|
||||
# 先进行混合检索，获取更多候选
|
||||
docs = self.hybrid_search(query, top_k * 3)
|
||||
|
||||
# 应用元数据过滤
|
||||
filtered_docs = []
|
||||
for doc in docs:
|
||||
match = True
|
||||
for key, value in filters.items():
|
||||
if key in doc.metadata:
|
||||
if isinstance(value, list):
|
||||
if doc.metadata[key] not in value:
|
||||
match = False
|
||||
break
|
||||
else:
|
||||
if doc.metadata[key] != value:
|
||||
match = False
|
||||
break
|
||||
else:
|
||||
match = False
|
||||
break
|
||||
|
||||
if match:
|
||||
filtered_docs.append(doc)
|
||||
if len(filtered_docs) >= top_k:
|
||||
break
|
||||
|
||||
return filtered_docs
|
||||
|
||||
def _rrf_rerank(self, vector_docs: List[Document], bm25_docs: List[Document], k: int = 60) -> List[Document]:
|
||||
"""
|
||||
使用RRF (Reciprocal Rank Fusion) 算法重排文档
|
||||
|
||||
Args:
|
||||
vector_docs: 向量检索结果
|
||||
bm25_docs: BM25检索结果
|
||||
k: RRF参数，用于平滑排名
|
||||
|
||||
Returns:
|
||||
重排后的文档列表
|
||||
"""
|
||||
doc_scores = {}
|
||||
doc_objects = {}
|
||||
|
||||
# 计算向量检索结果的RRF分数
|
||||
for rank, doc in enumerate(vector_docs):
|
||||
# 使用文档内容的哈希作为唯一标识
|
||||
doc_id = hash(doc.page_content)
|
||||
doc_objects[doc_id] = doc
|
||||
|
||||
# RRF公式: 1 / (k + rank)
|
||||
rrf_score = 1.0 / (k + rank + 1)
|
||||
doc_scores[doc_id] = doc_scores.get(doc_id, 0) + rrf_score
|
||||
|
||||
logger.debug(f"向量检索 - 文档{rank+1}: RRF分数 = {rrf_score:.4f}")
|
||||
|
||||
# 计算BM25检索结果的RRF分数
|
||||
for rank, doc in enumerate(bm25_docs):
|
||||
doc_id = hash(doc.page_content)
|
||||
doc_objects[doc_id] = doc
|
||||
|
||||
rrf_score = 1.0 / (k + rank + 1)
|
||||
doc_scores[doc_id] = doc_scores.get(doc_id, 0) + rrf_score
|
||||
|
||||
logger.debug(f"BM25检索 - 文档{rank+1}: RRF分数 = {rrf_score:.4f}")
|
||||
|
||||
# 按最终RRF分数排序
|
||||
sorted_docs = sorted(doc_scores.items(), key=lambda x: x[1], reverse=True)
|
||||
|
||||
# 构建最终结果
|
||||
reranked_docs = []
|
||||
for doc_id, final_score in sorted_docs:
|
||||
if doc_id in doc_objects:
|
||||
doc = doc_objects[doc_id]
|
||||
# 将RRF分数添加到文档元数据中
|
||||
doc.metadata['rrf_score'] = final_score
|
||||
reranked_docs.append(doc)
|
||||
logger.debug(f"最终排序 - 文档: {doc.page_content[:50]}... 最终RRF分数: {final_score:.4f}")
|
||||
|
||||
logger.info(f"RRF重排完成: 向量检索{len(vector_docs)}个文档, BM25检索{len(bm25_docs)}个文档, 合并后{len(reranked_docs)}个文档")
|
||||
|
||||
return reranked_docs
|
||||
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
langchain==0.3.26
|
||||
langchain-huggingface==0.3.1
|
||||
langchain-text-splitters==0.3.8
|
||||
langchain-unstructured==0.1.6
|
||||
langchain-community==0.3.27
|
||||
faiss-cpu>=1.7.0
|
||||
unstructured==0.18.11
|
||||
Markdown==3.8.2
|
||||
sentence-transformers>=3.0.0
|
||||
lazy_loader==0.4
|
||||
rank_bm25==0.2.2
|
||||
openai>=1.86.0,<2.0.0
|
||||
Binary file not shown.
@@ -0,0 +1,438 @@
|
||||
# AI菜谱知识图谱生成器
|
||||
## 基于Kimi大模型的智能菜谱解析系统
|
||||
|
||||
### 🌟 新功能特性
|
||||
|
||||
- **🤖 AI智能解析**: 使用Kimi大模型准确提取菜谱信息
|
||||
- **📊 结构化输出**: 自动生成标准化的知识图谱数据
|
||||
- **🔄 批量处理**: 支持大规模菜谱目录的批量转换
|
||||
- **💾 多格式导出**: 支持Neo4j和CSV两种输出格式
|
||||
- **🎯 高精度**: AI模型确保食材分类、步骤解析的准确性
|
||||
- **📁 智能目录识别**: 自动扫描dishes/目录，根据子目录名推断分类
|
||||
- **⚡ 优化处理**: 减少AI分类工作，提高处理效率和准确性
|
||||
|
||||
### 🚀 快速开始
|
||||
|
||||
#### 1. 安装依赖
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
#### 2. 配置API密钥
|
||||
|
||||
方法一：设置环境变量
|
||||
```bash
|
||||
export KIMI_API_KEY="your_api_key_here"
|
||||
```
|
||||
|
||||
方法二：编辑config.json
|
||||
```json
|
||||
{
|
||||
"kimi": {
|
||||
"api_key": "your_api_key_here"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### 3. 测试AI解析功能
|
||||
|
||||
```bash
|
||||
python run_ai_agent.py test
|
||||
```
|
||||
|
||||
#### 4. 批量处理菜谱
|
||||
|
||||
```bash
|
||||
# 处理HowToCook项目（会自动扫描dishes/目录）
|
||||
python run_ai_agent.py /path/to/HowToCook-master
|
||||
|
||||
# 处理其他菜谱目录
|
||||
python run_ai_agent.py /path/to/your/recipes
|
||||
```
|
||||
|
||||
### 📁 项目结构
|
||||
|
||||
```
|
||||
cook-rag-example/
|
||||
├── recipe_ai_agent.py # AI解析核心引擎
|
||||
├── run_ai_agent.py # 简化运行脚本
|
||||
├── batch_manager.py # 批次管理工具
|
||||
├── amount_normalizer.py # 用量标准化工具
|
||||
├── config.json # 配置文件
|
||||
├── requirements.txt # 依赖列表
|
||||
└── ai_output/ # AI输出目录
|
||||
├── nodes.csv # Neo4j节点数据
|
||||
├── relationships.csv # Neo4j关系数据
|
||||
└── neo4j_import.cypher # Neo4j导入脚本
|
||||
```
|
||||
|
||||
### 📁 智能目录处理
|
||||
|
||||
#### 自动分类识别
|
||||
|
||||
系统会智能识别HowToCook项目的目录结构：
|
||||
|
||||
- **专门扫描**: 只处理 `dishes/` 目录下的菜谱文件
|
||||
- **自动分类**: 根据子目录名自动推断菜谱分类
|
||||
- **排除无关文件**: 自动跳过 `template/`、`.github/` 等目录
|
||||
|
||||
#### 目录分类映射
|
||||
|
||||
```
|
||||
vegetable_dish/ → 素菜
|
||||
meat_dish/ → 荤菜
|
||||
aquatic/ → 水产
|
||||
breakfast/ → 早餐
|
||||
staple/ → 主食
|
||||
soup/ → 汤类
|
||||
dessert/ → 甜品
|
||||
drink/ → 饮料
|
||||
condiment/ → 调料
|
||||
semi-finished/ → 半成品
|
||||
```
|
||||
|
||||
#### 处理优化
|
||||
|
||||
- **减少AI工作量**: 分类信息直接从目录结构获取，无需AI推理
|
||||
- **提高准确性**: 避免AI分类错误，确保分类一致性
|
||||
- **加快处理速度**: 减少API调用复杂度和处理时间
|
||||
|
||||
### 🤖 AI解析能力
|
||||
|
||||
#### 智能信息提取
|
||||
|
||||
AI系统能够从Markdown菜谱中准确提取：
|
||||
|
||||
1. **基本信息**
|
||||
- 菜谱名称
|
||||
- 难度等级（1-5星）
|
||||
- 菜谱分类（素菜/荤菜/水产等）
|
||||
- 菜系归属（川菜/粤菜等）
|
||||
|
||||
2. **食材信息**
|
||||
- 食材名称和分类
|
||||
- 用量和单位
|
||||
- 主要食材识别
|
||||
|
||||
3. **烹饪步骤**
|
||||
- 步骤描述和顺序
|
||||
- 使用的烹饪方法
|
||||
- 需要的工具
|
||||
- 时间估计
|
||||
|
||||
4. **额外信息**
|
||||
- 准备时间/烹饪时间
|
||||
- 供应人数
|
||||
- 营养信息（当可用时）
|
||||
- 相关标签
|
||||
|
||||
#### 示例：AI解析结果
|
||||
|
||||
输入菜谱：
|
||||
```markdown
|
||||
# 红烧茄子的做法
|
||||
预估烹饪难度：★★★★
|
||||
## 必备原料和工具
|
||||
- 青茄子
|
||||
- 大蒜
|
||||
- 酱油
|
||||
- 面粉
|
||||
```
|
||||
|
||||
AI解析输出：
|
||||
```json
|
||||
{
|
||||
"name": "红烧茄子",
|
||||
"difficulty": 4,
|
||||
"category": "素菜",
|
||||
"ingredients": [
|
||||
{
|
||||
"name": "青茄子",
|
||||
"category": "蔬菜",
|
||||
"is_main": true
|
||||
},
|
||||
{
|
||||
"name": "大蒜",
|
||||
"category": "蔬菜",
|
||||
"is_main": false
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
### 📊 知识图谱结构
|
||||
|
||||
#### Neo4j数据模型
|
||||
|
||||
**节点类型**:
|
||||
- `Recipe`: 菜谱
|
||||
- `Ingredient`: 食材
|
||||
- `CookingStep`: 烹饪步骤
|
||||
|
||||
**关系类型**:
|
||||
- `has_ingredient`: 菜谱包含食材
|
||||
- `has_step`: 菜谱包含步骤
|
||||
- `belongs_to_category`: 属于分类
|
||||
- `has_difficulty`: 具有难度
|
||||
|
||||
#### 示例查询
|
||||
|
||||
```cypher
|
||||
// 查找所有包含茄子的菜谱
|
||||
MATCH (recipe:Concept)-[:has_ingredient]->(ing:Concept)
|
||||
WHERE ing.name CONTAINS "茄子"
|
||||
RETURN recipe.name, recipe.difficulty
|
||||
|
||||
// 查找四星难度的素菜
|
||||
MATCH (recipe:Concept)-[:belongs_to_category]->(cat:Concept)
|
||||
WHERE cat.name = "素菜" AND recipe.difficulty = 4
|
||||
RETURN recipe.name
|
||||
|
||||
// 推荐基于现有食材的菜谱
|
||||
MATCH (recipe:Concept)-[:has_ingredient]->(ing:Concept)
|
||||
WHERE ing.name IN ["茄子", "大蒜", "酱油"]
|
||||
WITH recipe, count(ing) as matches
|
||||
WHERE matches >= 2
|
||||
RETURN recipe.name, matches
|
||||
ORDER BY matches DESC
|
||||
```
|
||||
|
||||
### ⚙️ 配置选项
|
||||
|
||||
#### config.json详细配置
|
||||
|
||||
```json
|
||||
{
|
||||
"deepseek": {
|
||||
"api_key": "your_api_key",
|
||||
"base_url": "https://api.deepseek.com",
|
||||
"model": "deepseek-chat",
|
||||
"max_retries": 3,
|
||||
"timeout": 30
|
||||
},
|
||||
"processing": {
|
||||
"batch_size": 10,
|
||||
"delay_between_requests": 1,
|
||||
"max_concurrent_requests": 5
|
||||
},
|
||||
"output": {
|
||||
"format": "neo4j",
|
||||
"directory": "./ai_output",
|
||||
"include_nutrition": true,
|
||||
"include_tags": true
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### �️ 辅助工具
|
||||
|
||||
#### 批次管理工具 (batch_manager.py)
|
||||
|
||||
用于管理分批处理的进度和数据：
|
||||
|
||||
```bash
|
||||
# 查看处理状态
|
||||
python batch_manager.py status
|
||||
|
||||
# 继续中断的处理
|
||||
python batch_manager.py continue /path/to/recipes
|
||||
|
||||
# 合并批次数据
|
||||
python batch_manager.py merge
|
||||
|
||||
# 清理进度文件
|
||||
python batch_manager.py clean-progress
|
||||
|
||||
# 清理批次数据
|
||||
python batch_manager.py clean-batches
|
||||
|
||||
# 显示批次详情
|
||||
python batch_manager.py details
|
||||
```
|
||||
|
||||
#### 用量标准化工具 (amount_normalizer.py)
|
||||
|
||||
提供食材用量的标准化处理：
|
||||
|
||||
```python
|
||||
from amount_normalizer import AmountNormalizer
|
||||
|
||||
normalizer = AmountNormalizer()
|
||||
|
||||
# 标准化用量
|
||||
normalized, estimated = normalizer.normalize_amount("适量", "毫升")
|
||||
# 返回: ("适量", 10.0)
|
||||
|
||||
# 获取可比较的数值
|
||||
comparable = normalizer.get_comparable_value("一把", "")
|
||||
# 返回: 50.0
|
||||
|
||||
# 格式化显示
|
||||
display = normalizer.format_for_display("2-3个", "")
|
||||
# 返回: "2-3个"
|
||||
```
|
||||
|
||||
### �🔧 高级用法
|
||||
|
||||
#### 1. 自定义分类映射
|
||||
|
||||
```python
|
||||
# 在recipe_ai_agent.py中修改
|
||||
category_mapping = {
|
||||
"素菜": "710000000",
|
||||
"荤菜": "720000000",
|
||||
"自定义分类": "999000000"
|
||||
}
|
||||
```
|
||||
|
||||
#### 2. 扩展AI提示词
|
||||
|
||||
```python
|
||||
# 修改extract_recipe_info方法中的prompt
|
||||
prompt = f"""
|
||||
请分析菜谱并按以下格式提取信息：
|
||||
- 添加您的自定义要求
|
||||
- 特殊的分类规则
|
||||
- 额外的营养信息要求
|
||||
"""
|
||||
```
|
||||
|
||||
#### 3. 批量处理优化
|
||||
|
||||
```python
|
||||
# 调整处理参数
|
||||
builder = RecipeKnowledgeGraphBuilder(ai_agent)
|
||||
builder.batch_size = 20 # 增加批次大小
|
||||
builder.delay = 0.5 # 减少请求间隔
|
||||
```
|
||||
|
||||
### 📈 性能优化
|
||||
|
||||
#### API调用优化
|
||||
|
||||
1. **合理的请求频率**: 默认每秒1次请求，避免API限制
|
||||
2. **错误重试机制**: 自动重试失败的请求
|
||||
3. **批量处理**: 分批处理大量菜谱文件
|
||||
|
||||
#### 内存优化
|
||||
|
||||
1. **流式处理**: 逐个处理菜谱文件，避免内存溢出
|
||||
2. **定期清理**: 处理完成后及时释放内存
|
||||
3. **进度监控**: 实时显示处理进度
|
||||
|
||||
### 🔍 故障排除
|
||||
|
||||
#### 常见问题
|
||||
|
||||
1. **API密钥错误**
|
||||
```
|
||||
错误: API调用失败: 401
|
||||
解决: 检查API密钥是否正确设置
|
||||
```
|
||||
|
||||
2. **网络连接问题**
|
||||
```
|
||||
错误: API调用超时
|
||||
解决: 检查网络连接，或增加timeout设置
|
||||
```
|
||||
|
||||
3. **JSON解析错误**
|
||||
```
|
||||
错误: JSON解析错误
|
||||
解决: AI响应格式异常，会自动使用备用解析方法
|
||||
```
|
||||
|
||||
4. **菜谱格式问题**
|
||||
```
|
||||
错误: 无法提取菜谱信息
|
||||
解决: 检查Markdown格式是否符合要求
|
||||
```
|
||||
|
||||
#### 调试模式
|
||||
|
||||
```bash
|
||||
# 启用详细日志
|
||||
export DEBUG=true
|
||||
python run_ai_agent.py /path/to/recipes
|
||||
|
||||
# 测试单个菜谱
|
||||
python run_ai_agent.py test
|
||||
```
|
||||
|
||||
### 📋 使用场景
|
||||
|
||||
#### 1. 菜谱网站构建
|
||||
- 自动分类和标签
|
||||
- 智能推荐系统
|
||||
- 营养分析
|
||||
|
||||
#### 2. 烹饪应用开发
|
||||
- 食材识别
|
||||
- 步骤指导
|
||||
- 工具推荐
|
||||
|
||||
#### 3. 营养研究
|
||||
- 食材营养分析
|
||||
- 膳食搭配研究
|
||||
- 健康饮食推荐
|
||||
|
||||
#### 4. 餐饮业务
|
||||
- 菜单优化
|
||||
- 成本分析
|
||||
- 客户偏好分析
|
||||
|
||||
### 🌐 扩展开发
|
||||
|
||||
#### 添加新的AI模型
|
||||
|
||||
```python
|
||||
class CustomAIAgent(DeepSeekRecipeAgent):
|
||||
def __init__(self, api_key, model_name="custom-model"):
|
||||
super().__init__(api_key)
|
||||
self.model_name = model_name
|
||||
|
||||
def call_custom_api(self, messages):
|
||||
# 实现您的自定义AI调用逻辑
|
||||
pass
|
||||
```
|
||||
|
||||
#### 支持新的输出格式
|
||||
|
||||
```python
|
||||
def export_to_custom_format(self, output_dir):
|
||||
"""导出为自定义格式"""
|
||||
# 实现您的导出逻辑
|
||||
pass
|
||||
```
|
||||
|
||||
### 📊 数据质量保证
|
||||
|
||||
#### AI解析准确性
|
||||
|
||||
- **多轮验证**: AI提取后进行格式验证
|
||||
- **备用解析**: AI失败时使用规则解析
|
||||
- **人工审核**: 提供数据审核接口
|
||||
|
||||
#### 数据一致性
|
||||
|
||||
- **标准化分类**: 统一的食材和菜谱分类
|
||||
- **关系验证**: 确保图谱关系的逻辑一致性
|
||||
- **重复检测**: 自动识别和处理重复数据
|
||||
|
||||
---
|
||||
|
||||
**享受AI驱动的菜谱知识图谱构建体验！** 🎉
|
||||
|
||||
现在你可以安全地进行批量处理了！使用以下命令处理整个HowToCook项目：
|
||||
|
||||
```bash
|
||||
python run_ai_agent.py HowToCook-master
|
||||
```
|
||||
|
||||
系统会：
|
||||
1. 🎯 专门扫描 `HowToCook-master/dishes/` 目录
|
||||
2. 📁 根据子目录自动识别分类 (vegetable_dish→素菜, meat_dish→荤菜等)
|
||||
3. 🚫 自动排除template、.github等非菜谱目录
|
||||
4. 🤖 用AI智能解析每个菜谱的详细信息
|
||||
5. 📊 生成完整的知识图谱数据
|
||||
@@ -0,0 +1,222 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
用量标准化工具 - 处理中文菜谱中的非数值用量表达
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import Tuple, Optional
|
||||
|
||||
class AmountNormalizer:
|
||||
"""用量标准化器"""
|
||||
|
||||
def __init__(self):
|
||||
# 标准化映射表
|
||||
self.amount_mappings = {
|
||||
# 适量类
|
||||
"适量": "适量",
|
||||
"适当": "适量",
|
||||
"酌量": "适量",
|
||||
"随意": "适量",
|
||||
"看个人喜好": "适量",
|
||||
|
||||
# 少量类
|
||||
"少许": "少许",
|
||||
"少量": "少许",
|
||||
"一点点": "少许",
|
||||
"微量": "少许",
|
||||
"稍许": "少许",
|
||||
"稍微": "少许",
|
||||
|
||||
# 中等量类
|
||||
"中量": "中量",
|
||||
"适中": "中量",
|
||||
"正常": "中量",
|
||||
|
||||
# 大量类
|
||||
"大量": "大量",
|
||||
"足量": "大量",
|
||||
"充足": "大量",
|
||||
|
||||
# 数量词类
|
||||
"一把": "1把",
|
||||
"小把": "小把",
|
||||
"大把": "大把",
|
||||
"一勺": "1勺",
|
||||
"一小勺": "1小勺",
|
||||
"一大勺": "1大勺",
|
||||
"一茶匙": "1茶匙",
|
||||
"一汤匙": "1汤匙",
|
||||
"一撮": "1撮",
|
||||
|
||||
# 滴类
|
||||
"几滴": "几滴",
|
||||
"数滴": "几滴",
|
||||
"2-3滴": "几滴",
|
||||
|
||||
# 片/个/根类
|
||||
"几片": "几片",
|
||||
"数片": "几片",
|
||||
"几个": "几个",
|
||||
"数个": "几个",
|
||||
"几根": "几根",
|
||||
"数根": "几根",
|
||||
"几颗": "几颗",
|
||||
"数颗": "几颗"
|
||||
}
|
||||
|
||||
# 估算数值映射（用于营养计算等）
|
||||
self.estimated_values = {
|
||||
"适量": None, # 无法估算
|
||||
"少许": 2, # 约2g/ml
|
||||
"中量": 10, # 约10g/ml
|
||||
"大量": 30, # 约30g/ml
|
||||
"1把": 15, # 约15g
|
||||
"小把": 10, # 约10g
|
||||
"大把": 20, # 约20g
|
||||
"1勺": 5, # 约5ml
|
||||
"1小勺": 3, # 约3ml
|
||||
"1大勺": 15, # 约15ml
|
||||
"1茶匙": 5, # 约5ml
|
||||
"1汤匙": 15, # 约15ml
|
||||
"1撮": 1, # 约1g
|
||||
"几滴": 1, # 约1ml
|
||||
"几片": 3, # 约3片
|
||||
"几个": 3, # 约3个
|
||||
"几根": 3, # 约3根
|
||||
"几颗": 3 # 约3颗
|
||||
}
|
||||
|
||||
def normalize_amount(self, amount: str, unit: str = "") -> Tuple[str, Optional[float]]:
|
||||
"""
|
||||
标准化用量表达
|
||||
|
||||
Args:
|
||||
amount: 原始用量
|
||||
unit: 单位
|
||||
|
||||
Returns:
|
||||
Tuple[标准化后的用量, 估算数值]
|
||||
"""
|
||||
if not amount:
|
||||
return "", None
|
||||
|
||||
# 清理输入
|
||||
amount = amount.strip()
|
||||
|
||||
# 尝试提取数字
|
||||
number_match = re.match(r'^(\d+(?:\.\d+)?)', amount)
|
||||
if number_match:
|
||||
# 如果是纯数字，直接返回
|
||||
try:
|
||||
numeric_value = float(number_match.group(1))
|
||||
return amount, numeric_value
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# 标准化非数值表达
|
||||
normalized = self.amount_mappings.get(amount, amount)
|
||||
estimated = self.estimated_values.get(normalized, None)
|
||||
|
||||
return normalized, estimated
|
||||
|
||||
def parse_amount_with_unit(self, amount_str: str) -> Tuple[str, str, Optional[float]]:
|
||||
"""
|
||||
解析包含单位的用量字符串
|
||||
|
||||
Args:
|
||||
amount_str: 如"300毫升"、"适量盐"等
|
||||
|
||||
Returns:
|
||||
Tuple[数量, 单位, 估算数值]
|
||||
"""
|
||||
if not amount_str:
|
||||
return "", "", None
|
||||
|
||||
# 常见单位模式
|
||||
unit_patterns = [
|
||||
r'(\d+(?:\.\d+)?)\s*(克|千克|斤|两|钱)', # 重量
|
||||
r'(\d+(?:\.\d+)?)\s*(毫升|升|杯|勺|茶匙|汤匙)', # 体积
|
||||
r'(\d+(?:\.\d+)?)\s*(个|只|条|片|根|瓣|颗)', # 数量
|
||||
r'(\d+(?:\.\d+)?)\s*(把|撮|滴)', # 特殊
|
||||
]
|
||||
|
||||
# 尝试匹配数字+单位
|
||||
for pattern in unit_patterns:
|
||||
match = re.search(pattern, amount_str)
|
||||
if match:
|
||||
amount = match.group(1)
|
||||
unit = match.group(2)
|
||||
try:
|
||||
numeric_value = float(amount)
|
||||
return amount, unit, numeric_value
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# 处理纯文字表达
|
||||
normalized, estimated = self.normalize_amount(amount_str)
|
||||
return normalized, "", estimated
|
||||
|
||||
def get_comparable_value(self, amount: str, unit: str = "") -> Optional[float]:
|
||||
"""
|
||||
获取可比较的数值（用于排序、筛选等）
|
||||
|
||||
Args:
|
||||
amount: 用量
|
||||
unit: 单位
|
||||
|
||||
Returns:
|
||||
可比较的数值，None表示无法比较
|
||||
"""
|
||||
# 尝试解析为数字
|
||||
try:
|
||||
return float(amount)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# 使用估算值
|
||||
normalized, estimated = self.normalize_amount(amount, unit)
|
||||
return estimated
|
||||
|
||||
def format_for_display(self, amount: str, unit: str = "") -> str:
|
||||
"""
|
||||
格式化用于显示的用量字符串
|
||||
|
||||
Args:
|
||||
amount: 用量
|
||||
unit: 单位
|
||||
|
||||
Returns:
|
||||
格式化后的字符串
|
||||
"""
|
||||
normalized, _ = self.normalize_amount(amount, unit)
|
||||
|
||||
if unit and normalized:
|
||||
# 如果有单位且不是特殊表达（如"适量"已经很完整）
|
||||
if normalized not in ["适量", "少许", "中量", "大量"]:
|
||||
return f"{normalized} {unit}"
|
||||
else:
|
||||
return normalized
|
||||
else:
|
||||
return normalized or amount
|
||||
|
||||
# 使用示例
|
||||
def demo_normalization():
|
||||
"""演示标准化功能"""
|
||||
normalizer = AmountNormalizer()
|
||||
|
||||
test_cases = [
|
||||
("适量", "毫升"),
|
||||
("少许", "克"),
|
||||
("一把", ""),
|
||||
("300", "毫升"),
|
||||
("几滴", ""),
|
||||
("酌量", ""),
|
||||
("2-3滴", ""),
|
||||
("一小勺", ""),
|
||||
]
|
||||
|
||||
pass
|
||||
|
||||
if __name__ == "__main__":
|
||||
demo_normalization()
|
||||
@@ -0,0 +1,269 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
分批处理管理器 - 管理菜谱处理的分批保存和断点续传功能
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import sys
|
||||
import argparse
|
||||
from datetime import datetime
|
||||
from recipe_ai_agent import KimiRecipeAgent, RecipeKnowledgeGraphBuilder
|
||||
|
||||
def load_config():
|
||||
"""加载配置文件"""
|
||||
config_file = "config.json"
|
||||
if os.path.exists(config_file):
|
||||
with open(config_file, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
else:
|
||||
print("❌ 未找到config.json配置文件")
|
||||
sys.exit(1)
|
||||
|
||||
def show_progress_status(output_dir: str):
|
||||
"""显示处理进度状态"""
|
||||
progress_file = os.path.join(output_dir, "progress.json")
|
||||
|
||||
if not os.path.exists(progress_file):
|
||||
print("📋 未找到进度文件，没有正在进行的任务")
|
||||
return
|
||||
|
||||
try:
|
||||
with open(progress_file, 'r', encoding='utf-8') as f:
|
||||
progress_data = json.load(f)
|
||||
|
||||
print("处理进度状态:")
|
||||
print(f"总文件: {progress_data.get('total_files', 'N/A')}, 已处理: {progress_data.get('processed_count', 0)}")
|
||||
|
||||
status = progress_data.get('current_file', 'N/A')
|
||||
if status == 'COMPLETED':
|
||||
print("状态: 已完成")
|
||||
elif status == 'INTERRUPTED':
|
||||
print("状态: 被中断")
|
||||
else:
|
||||
print("状态: 进行中")
|
||||
|
||||
# 检查批次文件
|
||||
batch_dirs = [d for d in os.listdir(output_dir)
|
||||
if d.startswith("batch_") and os.path.isdir(os.path.join(output_dir, d))]
|
||||
|
||||
if batch_dirs:
|
||||
print(f"已保存批次: {len(batch_dirs)} 个")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ 读取进度文件失败: {str(e)}")
|
||||
|
||||
def clean_progress(output_dir: str):
|
||||
"""清理进度文件，重新开始处理"""
|
||||
progress_file = os.path.join(output_dir, "progress.json")
|
||||
|
||||
if os.path.exists(progress_file):
|
||||
confirm = input("⚠️ 确认要清理进度文件吗？这将删除所有处理进度 (y/N): ").strip().lower()
|
||||
if confirm == 'y':
|
||||
os.remove(progress_file)
|
||||
print("✅ 进度文件已清理")
|
||||
else:
|
||||
print("取消操作")
|
||||
else:
|
||||
print("📋 未找到进度文件")
|
||||
|
||||
def clean_batches(output_dir: str):
|
||||
"""清理所有批次数据"""
|
||||
batch_dirs = [d for d in os.listdir(output_dir)
|
||||
if d.startswith("batch_") and os.path.isdir(os.path.join(output_dir, d))]
|
||||
|
||||
if not batch_dirs:
|
||||
print("📁 未找到批次数据")
|
||||
return
|
||||
|
||||
print(f"找到 {len(batch_dirs)} 个批次目录:")
|
||||
for batch_dir in sorted(batch_dirs):
|
||||
print(f" - {batch_dir}")
|
||||
|
||||
confirm = input("\n⚠️ 确认要删除所有批次数据吗？ (y/N): ").strip().lower()
|
||||
if confirm == 'y':
|
||||
import shutil
|
||||
for batch_dir in batch_dirs:
|
||||
batch_path = os.path.join(output_dir, batch_dir)
|
||||
shutil.rmtree(batch_path)
|
||||
print(f"🗑️ 已删除: {batch_dir}")
|
||||
print("✅ 所有批次数据已清理")
|
||||
else:
|
||||
print("取消操作")
|
||||
|
||||
def merge_batches(output_dir: str):
|
||||
"""手动合并所有批次数据"""
|
||||
config = load_config()
|
||||
api_key = config["kimi"].get("api_key")
|
||||
|
||||
if not api_key:
|
||||
print("❌ 未找到API密钥配置")
|
||||
return
|
||||
|
||||
try:
|
||||
ai_agent = KimiRecipeAgent(api_key)
|
||||
builder = RecipeKnowledgeGraphBuilder(ai_agent, output_dir)
|
||||
|
||||
print("合并批次数据...")
|
||||
total_concepts, total_relationships = builder.merge_all_batches()
|
||||
|
||||
if total_concepts > 0 or total_relationships > 0:
|
||||
print(f"合并完成: {total_concepts} 概念, {total_relationships} 关系")
|
||||
|
||||
# 生成Neo4j格式
|
||||
format_type = config["output"].get("format", "neo4j")
|
||||
if format_type == "neo4j":
|
||||
builder.export_to_neo4j_csv(output_dir, merge_batches=False)
|
||||
else:
|
||||
print("未找到有效的批次数据")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ 合并失败: {str(e)}")
|
||||
|
||||
def continue_processing(recipe_dir: str, output_dir: str):
|
||||
"""继续处理中断的任务"""
|
||||
config = load_config()
|
||||
api_key = config["kimi"].get("api_key")
|
||||
|
||||
if not api_key:
|
||||
print("❌ 未找到API密钥配置")
|
||||
return
|
||||
|
||||
try:
|
||||
ai_agent = KimiRecipeAgent(api_key)
|
||||
batch_size = config.get("processing", {}).get("batch_size", 20)
|
||||
builder = RecipeKnowledgeGraphBuilder(ai_agent, output_dir, batch_size)
|
||||
|
||||
print("继续处理中断的任务...")
|
||||
processed, failed = builder.batch_process_recipes(recipe_dir, resume=True)
|
||||
|
||||
print(f"处理完成: 总数 {processed}, 失败 {failed}")
|
||||
|
||||
# 自动合并数据
|
||||
print("自动合并批次数据...")
|
||||
merge_batches(output_dir)
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ 处理失败: {str(e)}")
|
||||
|
||||
def show_batch_details(output_dir: str, batch_num: int = None):
|
||||
"""显示批次详细信息"""
|
||||
batch_dirs = [d for d in os.listdir(output_dir)
|
||||
if d.startswith("batch_") and os.path.isdir(os.path.join(output_dir, d))]
|
||||
|
||||
if not batch_dirs:
|
||||
print("📁 未找到批次数据")
|
||||
return
|
||||
|
||||
batch_dirs.sort()
|
||||
|
||||
if batch_num is not None:
|
||||
target_batch = f"batch_{batch_num:03d}"
|
||||
if target_batch not in batch_dirs:
|
||||
print(f"❌ 未找到批次 {batch_num}")
|
||||
return
|
||||
batch_dirs = [target_batch]
|
||||
|
||||
import pandas as pd
|
||||
|
||||
for batch_dir in batch_dirs:
|
||||
print(f"\n📁 {batch_dir}:")
|
||||
batch_path = os.path.join(output_dir, batch_dir)
|
||||
|
||||
# 概念文件
|
||||
concepts_file = os.path.join(batch_path, "concepts.csv")
|
||||
if os.path.exists(concepts_file):
|
||||
df = pd.read_csv(concepts_file)
|
||||
print(f" 概念数量: {len(df)}")
|
||||
|
||||
# 按类型统计
|
||||
if 'concept_type' in df.columns:
|
||||
type_counts = df['concept_type'].value_counts()
|
||||
for concept_type, count in type_counts.items():
|
||||
print(f" - {concept_type}: {count}")
|
||||
|
||||
# 关系文件
|
||||
relationships_file = os.path.join(batch_path, "relationships.csv")
|
||||
if os.path.exists(relationships_file):
|
||||
df = pd.read_csv(relationships_file)
|
||||
print(f" 关系数量: {len(df)}")
|
||||
|
||||
# 按类型统计
|
||||
if 'relationship_type' in df.columns:
|
||||
type_counts = df['relationship_type'].value_counts()
|
||||
for rel_type, count in type_counts.items():
|
||||
print(f" - {rel_type}: {count}")
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
parser = argparse.ArgumentParser(
|
||||
description='分批处理管理器 - 管理菜谱处理的分批保存和断点续传',
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
使用示例:
|
||||
python batch_manager.py status # 查看处理状态
|
||||
python batch_manager.py continue ./HowToCook-master # 继续中断的处理
|
||||
python batch_manager.py merge # 合并批次数据
|
||||
python batch_manager.py clean-progress # 清理进度文件
|
||||
python batch_manager.py clean-batches # 清理批次数据
|
||||
python batch_manager.py details # 显示所有批次详情
|
||||
python batch_manager.py details -b 1 # 显示指定批次详情
|
||||
""")
|
||||
|
||||
subparsers = parser.add_subparsers(dest='command', help='可用命令')
|
||||
|
||||
# status命令
|
||||
subparsers.add_parser('status', help='查看处理进度状态')
|
||||
|
||||
# continue命令
|
||||
continue_parser = subparsers.add_parser('continue', help='继续中断的处理')
|
||||
continue_parser.add_argument('recipe_dir', help='菜谱目录路径')
|
||||
|
||||
# merge命令
|
||||
subparsers.add_parser('merge', help='合并所有批次数据')
|
||||
|
||||
# clean命令
|
||||
subparsers.add_parser('clean-progress', help='清理进度文件')
|
||||
subparsers.add_parser('clean-batches', help='清理所有批次数据')
|
||||
|
||||
# details命令
|
||||
details_parser = subparsers.add_parser('details', help='显示批次详细信息')
|
||||
details_parser.add_argument('-b', '--batch', type=int, help='指定批次编号')
|
||||
|
||||
# 全局参数
|
||||
parser.add_argument('-o', '--output', default='./ai_output', help='输出目录路径')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.command:
|
||||
parser.print_help()
|
||||
return
|
||||
|
||||
# 确保输出目录存在
|
||||
os.makedirs(args.output, exist_ok=True)
|
||||
|
||||
print("🛠️ 分批处理管理器")
|
||||
print("=" * 50)
|
||||
|
||||
try:
|
||||
if args.command == 'status':
|
||||
show_progress_status(args.output)
|
||||
elif args.command == 'continue':
|
||||
continue_processing(args.recipe_dir, args.output)
|
||||
elif args.command == 'merge':
|
||||
merge_batches(args.output)
|
||||
elif args.command == 'clean-progress':
|
||||
clean_progress(args.output)
|
||||
elif args.command == 'clean-batches':
|
||||
clean_batches(args.output)
|
||||
elif args.command == 'details':
|
||||
show_batch_details(args.output, args.batch)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print(f"\n\n⏹️ 操作被用户中断")
|
||||
except Exception as e:
|
||||
print(f"\n❌ 操作失败: {str(e)}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,39 @@
|
||||
{
|
||||
"kimi": {
|
||||
"api_key": "sk-xxx",
|
||||
"base_url": "https://api.moonshot.cn/v1",
|
||||
"model": "kimi-k2-0711-preview",
|
||||
"max_retries": 3,
|
||||
"timeout": 30
|
||||
},
|
||||
"processing": {
|
||||
"batch_size": 10,
|
||||
"delay_between_requests": 1,
|
||||
"max_concurrent_requests": 5
|
||||
},
|
||||
"output": {
|
||||
"format": "neo4j",
|
||||
"directory": "./ai_output",
|
||||
"include_nutrition": true,
|
||||
"include_tags": true,
|
||||
"available_formats": ["csv", "neo4j", "rf2"]
|
||||
},
|
||||
"categories": {
|
||||
"素菜": "710000000",
|
||||
"荤菜": "720000000",
|
||||
"水产": "730000000",
|
||||
"早餐": "740000000",
|
||||
"主食": "750000000",
|
||||
"汤类": "760000000",
|
||||
"甜品": "770000000",
|
||||
"饮料": "780000000",
|
||||
"调料": "790000000"
|
||||
},
|
||||
"difficulty_mapping": {
|
||||
"1": "610000000",
|
||||
"2": "620000000",
|
||||
"3": "630000000",
|
||||
"4": "640000000",
|
||||
"5": "650000000"
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,304 @@
|
||||
# 菜谱知识图谱设计方案
|
||||
## 烹饪本体设计
|
||||
|
||||
### 概述
|
||||
本设计方案将菜谱数据结构化为本体图数据库，支持语义查询、推理和知识发现。
|
||||
|
||||
### 核心概念体系 (Concept Hierarchy)
|
||||
|
||||
#### 1. 根概念 (Root Concept)
|
||||
```
|
||||
100000000 | Culinary Concept | (烹饪概念)
|
||||
```
|
||||
|
||||
#### 2. 顶级概念 (Top Level Concepts)
|
||||
```
|
||||
200000000 | Recipe | (菜谱)
|
||||
300000000 | Ingredient | (食材)
|
||||
400000000 | Tool | (工具)
|
||||
500000000 | Cooking Method | (烹饪方法)
|
||||
600000000 | Difficulty Level | (难度等级)
|
||||
700000000 | Recipe Category | (菜谱分类)
|
||||
800000000 | Measurement Unit | (计量单位)
|
||||
900000000 | Cooking Step | (烹饪步骤)
|
||||
```
|
||||
|
||||
### 概念分类体系
|
||||
|
||||
#### 菜谱分类 (Recipe Categories)
|
||||
```
|
||||
710000000 | Vegetable Dish | (素菜)
|
||||
720000000 | Meat Dish | (荤菜)
|
||||
730000000 | Aquatic Product | (水产)
|
||||
740000000 | Breakfast | (早餐)
|
||||
750000000 | Staple Food | (主食)
|
||||
760000000 | Soup | (汤类)
|
||||
770000000 | Dessert | (甜品)
|
||||
780000000 | Beverage | (饮料)
|
||||
790000000 | Condiment | (调料)
|
||||
```
|
||||
|
||||
#### 食材分类 (Ingredient Categories)
|
||||
```
|
||||
310000000 | Vegetable | (蔬菜)
|
||||
320000000 | Seasoning | (调料)
|
||||
330000000 | Protein | (蛋白质)
|
||||
340000000 | Starch | (淀粉类)
|
||||
350000000 | Dairy | (乳制品)
|
||||
360000000 | Fruit | (水果)
|
||||
370000000 | Herb | (香草香料)
|
||||
380000000 | Oil Fat | (油脂类)
|
||||
```
|
||||
|
||||
#### 难度等级 (Difficulty Levels)
|
||||
```
|
||||
610000000 | One Star | (一星) ★
|
||||
620000000 | Two Star | (二星) ★★
|
||||
630000000 | Three Star | (三星) ★★★
|
||||
640000000 | Four Star | (四星) ★★★★
|
||||
650000000 | Five Star | (五星) ★★★★★
|
||||
```
|
||||
|
||||
#### 烹饪方法 (Cooking Methods)
|
||||
```
|
||||
501000000 | Stir Fry | (炒)
|
||||
502000000 | Deep Fry | (炸)
|
||||
503000000 | Braise | (红烧)
|
||||
504000000 | Steam | (蒸)
|
||||
505000000 | Boil | (煮)
|
||||
506000000 | Roast | (烤)
|
||||
507000000 | Stew | (炖)
|
||||
508000000 | Mix | (拌)
|
||||
509000000 | Marinate | (腌)
|
||||
510000000 | Blanch | (焯)
|
||||
```
|
||||
|
||||
### 关系类型定义 (Relationship Types)
|
||||
|
||||
#### 核心关系 (Core Relationships)
|
||||
```
|
||||
116680003 | is_a | (是一个) - 层次分类关系
|
||||
```
|
||||
|
||||
#### 属性关系 (Attribute Relationships)
|
||||
```
|
||||
801000001 | has_ingredient | (包含食材)
|
||||
801000002 | requires_tool | (需要工具)
|
||||
801000003 | has_step | (包含步骤)
|
||||
801000004 | belongs_to_category | (属于分类)
|
||||
801000005 | has_difficulty | (具有难度)
|
||||
801000006 | uses_method | (使用方法)
|
||||
801000007 | has_amount | (具有用量)
|
||||
801000008 | step_follows | (步骤顺序)
|
||||
801000009 | serves_people | (供应人数)
|
||||
801000010 | cooking_time | (烹饪时间)
|
||||
801000011 | prep_time | (准备时间)
|
||||
801000012 | ingredient_substitute | (食材替代)
|
||||
801000013 | recipe_variant | (菜谱变体)
|
||||
801000014 | nutritional_info | (营养信息)
|
||||
```
|
||||
|
||||
### 具体实例设计
|
||||
|
||||
#### 红烧茄子菜谱实例
|
||||
```
|
||||
概念ID: 201000001
|
||||
完全限定名: 201000001 | 红烧茄子 (Braised Eggplant) |
|
||||
首选术语: 红烧茄子
|
||||
同义词: 茄子烧制, 红烧青茄子
|
||||
|
||||
属性关系:
|
||||
- belongs_to_category = 710000000 | 素菜
|
||||
- has_difficulty = 640000000 | 四星
|
||||
- serves_people = 2人份
|
||||
- cooking_time = 30分钟
|
||||
- prep_time = 15分钟
|
||||
|
||||
食材关系:
|
||||
- has_ingredient = 311000001 | 青茄子 | : has_amount = "0.7个/份"
|
||||
- has_ingredient = 311000002 | 大蒜 | : has_amount = "3瓣"
|
||||
- has_ingredient = 321000001 | 酱油 | : has_amount = "茄子数量*7克"
|
||||
- has_ingredient = 331000001 | 鸡蛋 | : has_amount = "1个"
|
||||
- has_ingredient = 341000001 | 面粉 | : has_amount = "青茄子数量*150克"
|
||||
|
||||
工具关系:
|
||||
- requires_tool = 401000001 | 炒锅
|
||||
- requires_tool = 401000002 | 菜刀
|
||||
- requires_tool = 401000003 | 筷子
|
||||
|
||||
方法关系:
|
||||
- uses_method = 502000000 | 炸
|
||||
- uses_method = 503000000 | 红烧
|
||||
- uses_method = 501000000 | 炒
|
||||
|
||||
步骤关系:
|
||||
- has_step = S001 | 清洗食材
|
||||
- has_step = S002 | 切配处理
|
||||
- has_step = S003 | 调制面糊
|
||||
- has_step = S004 | 油炸茄块
|
||||
- has_step = S005 | 炒制调味
|
||||
```
|
||||
|
||||
### 表达式系统
|
||||
|
||||
#### 预协调表达式 (Precoordinated)
|
||||
```
|
||||
201000001 | 红烧茄子 |
|
||||
```
|
||||
|
||||
#### 后协调表达式 (Postcoordinated)
|
||||
```
|
||||
# 四星难度的素菜
|
||||
710000000 | 素菜 | : has_difficulty = 640000000 | 四星 |
|
||||
|
||||
# 包含茄子的红烧菜谱
|
||||
200000000 | 菜谱 | : {
|
||||
has_ingredient = 311000001 | 青茄子 |,
|
||||
uses_method = 503000000 | 红烧 |
|
||||
}
|
||||
|
||||
# 30分钟内完成的四星菜谱
|
||||
200000000 | 菜谱 | : {
|
||||
has_difficulty = 640000000 | 四星 |,
|
||||
cooking_time <= 30分钟
|
||||
}
|
||||
```
|
||||
|
||||
### 数据文件结构
|
||||
|
||||
#### 概念文件 (rf2_concept.txt)
|
||||
```
|
||||
id effectiveTime active moduleId definitionStatusId
|
||||
100000000 20241201 1 900000000 900000000
|
||||
200000000 20241201 1 900000000 900000000
|
||||
201000001 20241201 1 900000000 900000000
|
||||
```
|
||||
|
||||
#### 描述文件 (rf2_description.txt)
|
||||
```
|
||||
id effectiveTime active moduleId conceptId languageCode typeId term caseSignificanceId
|
||||
D001 20241201 1 900000000 201000001 zh-CN 900000001 红烧茄子 900000000
|
||||
D002 20241201 1 900000000 201000001 zh-CN 900000002 茄子烧制 900000000
|
||||
D003 20241201 1 900000000 201000001 en 900000001 Braised Eggplant 900000000
|
||||
```
|
||||
|
||||
#### 关系文件 (rf2_relationship.txt)
|
||||
```
|
||||
id effectiveTime active moduleId sourceId destinationId relationshipGroup typeId characteristicTypeId modifierId
|
||||
R001 20241201 1 900000000 201000001 710000000 0 801000004 900000000 900000000
|
||||
R002 20241201 1 900000000 201000001 640000000 0 801000005 900000000 900000000
|
||||
R003 20241201 1 900000000 201000001 311000001 1 801000001 900000000 900000000
|
||||
```
|
||||
|
||||
### 查询示例
|
||||
|
||||
#### 1. 基础查询 - 所有素菜
|
||||
```cypher
|
||||
MATCH (recipe:Concept)-[:IS_A*]->(category:Concept {conceptId: "710000000"})
|
||||
RETURN recipe.preferredTerm
|
||||
```
|
||||
|
||||
#### 2. 复杂查询 - 包含特定食材的四星菜谱
|
||||
```cypher
|
||||
MATCH (recipe:Concept)-[:HAS_INGREDIENT]->(ingredient:Concept)
|
||||
WHERE ingredient.conceptId = "311000001"
|
||||
AND (recipe)-[:HAS_DIFFICULTY]->(:Concept {conceptId: "640000000"})
|
||||
RETURN recipe.preferredTerm, recipe.cookingTime
|
||||
```
|
||||
|
||||
#### 3. 语义查询 - 所有炒制类菜谱
|
||||
```cypher
|
||||
MATCH (recipe:Concept)-[:USES_METHOD]->(method:Concept)
|
||||
WHERE (method)-[:IS_A*]->(:Concept {conceptId: "501000000"})
|
||||
RETURN DISTINCT recipe.preferredTerm
|
||||
```
|
||||
|
||||
### 实现技术栈
|
||||
|
||||
#### 图数据库选择
|
||||
- **Neo4j**: 成熟的图数据库，适合复杂查询
|
||||
- **ArangoDB**: 多模型数据库，支持图和文档
|
||||
- **Amazon Neptune**: 云原生图数据库
|
||||
|
||||
#### API设计
|
||||
```python
|
||||
class RecipeOntologyAPI:
|
||||
def search_recipes_by_ingredient(self, ingredient_id: str) -> List[Recipe]:
|
||||
"""根据食材搜索菜谱"""
|
||||
pass
|
||||
|
||||
def get_recipe_variants(self, recipe_id: str) -> List[Recipe]:
|
||||
"""获取菜谱变体"""
|
||||
pass
|
||||
|
||||
def suggest_substitutes(self, ingredient_id: str) -> List[Ingredient]:
|
||||
"""建议食材替代"""
|
||||
pass
|
||||
|
||||
def analyze_nutrition(self, recipe_id: str) -> NutritionInfo:
|
||||
"""分析营养成分"""
|
||||
pass
|
||||
```
|
||||
|
||||
### 应用场景
|
||||
|
||||
#### 1. 智能菜谱推荐
|
||||
- 基于现有食材推荐菜谱
|
||||
- 根据难度等级筛选
|
||||
- 营养搭配建议
|
||||
|
||||
#### 2. 食材替代建议
|
||||
- 过敏原替代
|
||||
- 地域性食材替换
|
||||
- 营养等价替代
|
||||
|
||||
#### 3. 烹饪知识推理
|
||||
- 步骤优化建议
|
||||
- 工具使用指导
|
||||
- 时间管理优化
|
||||
|
||||
#### 4. 营养分析
|
||||
- 卡路里计算
|
||||
- 营养成分分析
|
||||
- 膳食搭配建议
|
||||
|
||||
### 扩展性设计
|
||||
|
||||
#### 多语言支持
|
||||
```
|
||||
zh-CN: 红烧茄子
|
||||
en-US: Braised Eggplant
|
||||
ja-JP: 茄子の煮物
|
||||
ko-KR: 가지조림
|
||||
```
|
||||
|
||||
#### 地域化扩展
|
||||
```
|
||||
中式烹饪: 701000000
|
||||
法式烹饪: 702000000
|
||||
意式烹饪: 703000000
|
||||
日式烹饪: 704000000
|
||||
```
|
||||
|
||||
#### 个性化标签
|
||||
```
|
||||
vegetarian: 素食主义
|
||||
vegan: 纯素
|
||||
halal: 清真
|
||||
kosher: 犹太教食物
|
||||
gluten_free: 无麸质
|
||||
```
|
||||
|
||||
### 质量控制
|
||||
|
||||
#### 概念一致性检查
|
||||
- 循环依赖检测
|
||||
- 概念完整性验证
|
||||
- 关系合理性校验
|
||||
|
||||
#### 数据质量保证
|
||||
- 重复概念检测
|
||||
- 缺失关系补充
|
||||
- 术语标准化
|
||||
|
||||
这个设计方案提供了一个完整的菜谱知识图谱框架，可以支持复杂的语义查询、推理和知识发现功能。
|
||||
@@ -0,0 +1,6 @@
|
||||
pandas==2.3.1
|
||||
pathlib2==2.3.6
|
||||
neo4j==5.28.1
|
||||
requests==2.32.4
|
||||
dataclasses==0.6
|
||||
openai==1.97.0
|
||||
@@ -0,0 +1,208 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
简化的AI菜谱解析运行脚本
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import sys
|
||||
from recipe_ai_agent import KimiRecipeAgent, RecipeKnowledgeGraphBuilder
|
||||
|
||||
def load_config():
|
||||
"""加载配置文件"""
|
||||
config_file = "config.json"
|
||||
if os.path.exists(config_file):
|
||||
with open(config_file, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
else:
|
||||
print("警告: 未找到config.json配置文件，将使用默认配置")
|
||||
return {
|
||||
"kimi": {
|
||||
"api_key": "",
|
||||
"base_url": "https://api.moonshot.cn/v1"
|
||||
},
|
||||
"output": {
|
||||
"format": "neo4j",
|
||||
"directory": "./ai_output"
|
||||
}
|
||||
}
|
||||
|
||||
def setup_api_key():
|
||||
"""设置API密钥"""
|
||||
api_key = os.getenv('KIMI_API_KEY')
|
||||
if not api_key:
|
||||
api_key = input("请输入Kimi API密钥: ").strip()
|
||||
if not api_key:
|
||||
print("错误: 必须提供API密钥")
|
||||
sys.exit(1)
|
||||
return api_key
|
||||
|
||||
def get_recipe_directory():
|
||||
"""获取菜谱目录"""
|
||||
if len(sys.argv) > 1:
|
||||
recipe_dir = sys.argv[1]
|
||||
else:
|
||||
recipe_dir = input("请输入菜谱目录路径: ").strip()
|
||||
|
||||
if not os.path.exists(recipe_dir):
|
||||
print(f"错误: 目录不存在 - {recipe_dir}")
|
||||
sys.exit(1)
|
||||
|
||||
return recipe_dir
|
||||
|
||||
def test_single_recipe():
|
||||
"""测试单个菜谱解析"""
|
||||
test_recipe = """# 红烧茄子的做法
|
||||
预估烹饪难度：★★★★
|
||||
|
||||
## 必备原料和工具
|
||||
- 大蒜
|
||||
- 大葱
|
||||
- 青辣椒
|
||||
- 洋葱
|
||||
- 西红柿
|
||||
- 青茄子
|
||||
- 盐
|
||||
- 酱油
|
||||
- 鸡蛋
|
||||
- 面粉
|
||||
- 淀粉
|
||||
|
||||
## 计算
|
||||
每次制作前需要确定计划做几份。一份正好够 2 个人食用
|
||||
|
||||
## 操作
|
||||
1. 青茄子、青辣椒、西红柿、洋葱、大葱洗净。
|
||||
2. 大葱切 5 毫米宽的葱花，大蒜扒皮并拍碎，西红柿切 6 立方厘米的块。
|
||||
3. 茄子切菱形块。
|
||||
4. 将面粉倒入盆中，依次加入少量水，搅拌均匀，呈粘稠糊状。
|
||||
5. 热锅，放入茄块翻炒至金黄色。
|
||||
"""
|
||||
|
||||
print("=== 测试单个菜谱解析 ===")
|
||||
|
||||
# 加载配置
|
||||
config = load_config()
|
||||
api_key = config["kimi"].get("api_key")
|
||||
if not api_key or api_key == "YOUR_KIMI_API_KEY_HERE":
|
||||
api_key = setup_api_key()
|
||||
|
||||
try:
|
||||
agent = KimiRecipeAgent(api_key)
|
||||
recipe_info = agent.extract_recipe_info(test_recipe, "dishes/vegetable_dish/红烧茄子.md")
|
||||
|
||||
print(f"测试成功: {recipe_info.name} ({len(recipe_info.ingredients)}个食材, {len(recipe_info.steps)}个步骤)")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ 测试失败: {str(e)}")
|
||||
return False
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
print("🍳 AI菜谱知识图谱生成器")
|
||||
print("=" * 50)
|
||||
|
||||
# 检查是否为测试模式
|
||||
if len(sys.argv) > 1 and sys.argv[1] == "test":
|
||||
success = test_single_recipe()
|
||||
sys.exit(0 if success else 1)
|
||||
|
||||
# 加载配置
|
||||
config = load_config()
|
||||
|
||||
# 设置API密钥
|
||||
api_key = config["kimi"].get("api_key")
|
||||
if not api_key or api_key == "YOUR_KIMI_API_KEY_HERE":
|
||||
api_key = setup_api_key()
|
||||
|
||||
# 获取菜谱目录
|
||||
recipe_dir = get_recipe_directory()
|
||||
|
||||
# 确认参数
|
||||
print(f"\n配置信息:")
|
||||
print(f"- API密钥: {api_key[:8]}...")
|
||||
print(f"- 菜谱目录: {recipe_dir}")
|
||||
print(f"- 输出格式: {config['output'].get('format', 'neo4j')}")
|
||||
print(f"- 输出目录: {config['output'].get('directory', './ai_output')}")
|
||||
|
||||
confirm = input("\n确认开始处理? (y/N): ").strip().lower()
|
||||
if confirm != 'y':
|
||||
print("取消处理")
|
||||
return
|
||||
|
||||
try:
|
||||
# 创建AI agent
|
||||
print("\n🤖 初始化AI Agent...")
|
||||
ai_agent = KimiRecipeAgent(api_key, config["kimi"].get("base_url"))
|
||||
|
||||
# 创建知识图谱构建器
|
||||
output_dir = config["output"].get("directory", "./ai_output")
|
||||
batch_size = config.get("processing", {}).get("batch_size", 20) # 默认批次大小为20
|
||||
builder = RecipeKnowledgeGraphBuilder(ai_agent, output_dir, batch_size)
|
||||
|
||||
# 批量处理菜谱
|
||||
print(f"\n📚 开始处理菜谱目录...")
|
||||
processed, failed = builder.batch_process_recipes(recipe_dir)
|
||||
|
||||
print(f"处理结果: 成功 {processed} 个，失败 {failed} 个")
|
||||
|
||||
# 导出数据
|
||||
output_dir = config["output"].get("directory", "./ai_output")
|
||||
output_format = config["output"].get("format", "neo4j")
|
||||
|
||||
print(f"导出数据 (格式: {output_format})...")
|
||||
|
||||
if output_format == "neo4j":
|
||||
builder.export_to_neo4j_csv(output_dir)
|
||||
print(f"Neo4j文件已生成: {output_dir}")
|
||||
elif output_format == "rf2":
|
||||
builder.export_to_rf2_format(output_dir)
|
||||
print(f"RF2文件已生成: {output_dir}")
|
||||
else:
|
||||
builder.export_to_csv(output_dir)
|
||||
print(f"CSV文件已生成: {output_dir}")
|
||||
|
||||
print("处理完成!")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print(f"\n\n⏹️ 用户中断处理")
|
||||
except Exception as e:
|
||||
print(f"\n❌ 处理过程中出现错误: {str(e)}")
|
||||
print(f"请检查API密钥、网络连接和菜谱文件格式")
|
||||
|
||||
def show_help():
|
||||
"""显示帮助信息"""
|
||||
help_text = """
|
||||
🍳 AI菜谱知识图谱生成器 - 使用指南
|
||||
|
||||
基本用法:
|
||||
python run_ai_agent.py [菜谱目录路径]
|
||||
|
||||
测试模式:
|
||||
python run_ai_agent.py test
|
||||
|
||||
环境变量:
|
||||
KIMI_API_KEY - Kimi API密钥
|
||||
|
||||
配置文件:
|
||||
config.json - 详细配置选项
|
||||
|
||||
示例:
|
||||
python run_ai_agent.py ./HowToCook-master
|
||||
python run_ai_agent.py test
|
||||
|
||||
输出格式:
|
||||
- neo4j: 生成Neo4j导入格式的CSV文件
|
||||
- csv: 生成标准CSV文件
|
||||
|
||||
更多信息请查看README.md
|
||||
"""
|
||||
print(help_text)
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) > 1 and sys.argv[1] in ["-h", "--help", "help"]:
|
||||
show_help()
|
||||
else:
|
||||
main()
|
||||
@@ -0,0 +1,73 @@
|
||||
"""
|
||||
基于图数据库的RAG系统配置文件
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, Any
|
||||
|
||||
@dataclass
|
||||
class GraphRAGConfig:
|
||||
"""基于图数据库的RAG系统配置类"""
|
||||
|
||||
# Neo4j数据库配置
|
||||
neo4j_uri: str = "bolt://localhost:7687"
|
||||
neo4j_user: str = "neo4j"
|
||||
neo4j_password: str = "all-in-rag"
|
||||
neo4j_database: str = "neo4j"
|
||||
|
||||
# Milvus配置
|
||||
milvus_host: str = "localhost"
|
||||
milvus_port: int = 19530
|
||||
milvus_collection_name: str = "cooking_knowledge"
|
||||
milvus_dimension: int = 512 # BGE-small-zh-v1.5的向量维度
|
||||
|
||||
# 模型配置
|
||||
embedding_model: str = "BAAI/bge-small-zh-v1.5"
|
||||
llm_model: str = "kimi-k2-0711-preview"
|
||||
|
||||
# 检索配置（LightRAG Round-robin策略）
|
||||
top_k: int = 5
|
||||
|
||||
# 生成配置
|
||||
temperature: float = 0.1
|
||||
max_tokens: int = 2048
|
||||
|
||||
# 图数据处理配置
|
||||
chunk_size: int = 500
|
||||
chunk_overlap: int = 50
|
||||
max_graph_depth: int = 2 # 图遍历最大深度
|
||||
|
||||
def __post_init__(self):
|
||||
"""初始化后的处理"""
|
||||
# LightRAG使用Round-robin策略，无需权重验证
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, config_dict: Dict[str, Any]) -> 'GraphRAGConfig':
|
||||
"""从字典创建配置对象"""
|
||||
return cls(**config_dict)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""转换为字典"""
|
||||
return {
|
||||
'neo4j_uri': self.neo4j_uri,
|
||||
'neo4j_user': self.neo4j_user,
|
||||
'neo4j_password': self.neo4j_password,
|
||||
'neo4j_database': self.neo4j_database,
|
||||
'milvus_host': self.milvus_host,
|
||||
'milvus_port': self.milvus_port,
|
||||
'milvus_collection_name': self.milvus_collection_name,
|
||||
'milvus_dimension': self.milvus_dimension,
|
||||
'embedding_model': self.embedding_model,
|
||||
'llm_model': self.llm_model,
|
||||
'top_k': self.top_k,
|
||||
|
||||
'temperature': self.temperature,
|
||||
'max_tokens': self.max_tokens,
|
||||
'chunk_size': self.chunk_size,
|
||||
'chunk_overlap': self.chunk_overlap,
|
||||
'max_graph_depth': self.max_graph_depth
|
||||
}
|
||||
|
||||
# 默认配置实例
|
||||
DEFAULT_CONFIG = GraphRAGConfig()
|
||||
+441
@@ -0,0 +1,441 @@
|
||||
"""
|
||||
基于图RAG的智能烹饪助手 - 主程序
|
||||
整合传统检索和图RAG检索，实现真正的图数据优势
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
# 设置日志
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 添加当前目录到Python路径
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from config import DEFAULT_CONFIG, GraphRAGConfig
|
||||
from rag_modules import (
|
||||
GraphDataPreparationModule,
|
||||
MilvusIndexConstructionModule,
|
||||
GenerationIntegrationModule
|
||||
)
|
||||
from rag_modules.hybrid_retrieval import HybridRetrievalModule
|
||||
from rag_modules.graph_rag_retrieval import GraphRAGRetrieval
|
||||
from rag_modules.intelligent_query_router import IntelligentQueryRouter, QueryAnalysis
|
||||
|
||||
# 加载环境变量
|
||||
load_dotenv()
|
||||
|
||||
class AdvancedGraphRAGSystem:
|
||||
"""
|
||||
图RAG系统
|
||||
|
||||
核心特性：
|
||||
1. 智能路由：自动选择最适合的检索策略
|
||||
2. 双引擎检索：传统混合检索 + 图RAG检索
|
||||
3. 图结构推理：多跳遍历、子图提取、关系推理
|
||||
4. 查询复杂度分析：深度理解用户意图
|
||||
5. 自适应学习：基于反馈优化系统性能
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[GraphRAGConfig] = None):
|
||||
self.config = config or DEFAULT_CONFIG
|
||||
|
||||
# 核心模块
|
||||
self.data_module = None
|
||||
self.index_module = None
|
||||
self.generation_module = None
|
||||
|
||||
# 检索引擎
|
||||
self.traditional_retrieval = None
|
||||
self.graph_rag_retrieval = None
|
||||
self.query_router = None
|
||||
|
||||
# 系统状态
|
||||
self.system_ready = False
|
||||
|
||||
def initialize_system(self):
|
||||
"""初始化高级图RAG系统"""
|
||||
logger.info("启动高级图RAG系统...")
|
||||
|
||||
try:
|
||||
# 1. 数据准备模块
|
||||
print("初始化数据准备模块...")
|
||||
self.data_module = GraphDataPreparationModule(
|
||||
uri=self.config.neo4j_uri,
|
||||
user=self.config.neo4j_user,
|
||||
password=self.config.neo4j_password,
|
||||
database=self.config.neo4j_database
|
||||
)
|
||||
|
||||
# 2. 向量索引模块
|
||||
print("初始化Milvus向量索引...")
|
||||
self.index_module = MilvusIndexConstructionModule(
|
||||
host=self.config.milvus_host,
|
||||
port=self.config.milvus_port,
|
||||
collection_name=self.config.milvus_collection_name,
|
||||
dimension=self.config.milvus_dimension,
|
||||
model_name=self.config.embedding_model
|
||||
)
|
||||
|
||||
# 3. 生成模块
|
||||
print("初始化生成模块...")
|
||||
self.generation_module = GenerationIntegrationModule(
|
||||
model_name=self.config.llm_model,
|
||||
temperature=self.config.temperature,
|
||||
max_tokens=self.config.max_tokens
|
||||
)
|
||||
|
||||
# 4. 传统混合检索模块
|
||||
print("初始化传统混合检索...")
|
||||
self.traditional_retrieval = HybridRetrievalModule(
|
||||
config=self.config,
|
||||
milvus_module=self.index_module,
|
||||
data_module=self.data_module,
|
||||
llm_client=self.generation_module.client
|
||||
)
|
||||
|
||||
# 5. 图RAG检索模块
|
||||
print("初始化图RAG检索引擎...")
|
||||
self.graph_rag_retrieval = GraphRAGRetrieval(
|
||||
config=self.config,
|
||||
llm_client=self.generation_module.client
|
||||
)
|
||||
|
||||
# 6. 智能查询路由器
|
||||
print("初始化智能查询路由器...")
|
||||
self.query_router = IntelligentQueryRouter(
|
||||
traditional_retrieval=self.traditional_retrieval,
|
||||
graph_rag_retrieval=self.graph_rag_retrieval,
|
||||
llm_client=self.generation_module.client,
|
||||
config=self.config
|
||||
)
|
||||
|
||||
print("✅ 高级图RAG系统初始化完成！")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"系统初始化失败: {e}")
|
||||
raise
|
||||
|
||||
def build_knowledge_base(self):
|
||||
"""构建知识库（如果需要）"""
|
||||
print("\n检查知识库状态...")
|
||||
|
||||
try:
|
||||
# 检查Milvus集合是否存在
|
||||
if self.index_module.has_collection():
|
||||
print("✅ 发现已存在的知识库，尝试加载...")
|
||||
if self.index_module.load_collection():
|
||||
print("知识库加载成功！")
|
||||
|
||||
# 重要：即使从已存在的知识库加载，也需要加载图数据以支持图索引
|
||||
print("加载图数据以支持图检索...")
|
||||
self.data_module.load_graph_data()
|
||||
print("构建菜谱文档...")
|
||||
self.data_module.build_recipe_documents()
|
||||
print("进行文档分块...")
|
||||
chunks = self.data_module.chunk_documents(
|
||||
chunk_size=self.config.chunk_size,
|
||||
chunk_overlap=self.config.chunk_overlap
|
||||
)
|
||||
|
||||
self._initialize_retrievers(chunks)
|
||||
return
|
||||
else:
|
||||
print("❌ 知识库加载失败，开始重建...")
|
||||
|
||||
print("未找到已存在的集合，开始构建新的知识库...")
|
||||
|
||||
# 从Neo4j加载图数据
|
||||
print("从Neo4j加载图数据...")
|
||||
self.data_module.load_graph_data()
|
||||
|
||||
# 构建菜谱文档
|
||||
print("构建菜谱文档...")
|
||||
self.data_module.build_recipe_documents()
|
||||
|
||||
# 进行文档分块
|
||||
print("进行文档分块...")
|
||||
chunks = self.data_module.chunk_documents(
|
||||
chunk_size=self.config.chunk_size,
|
||||
chunk_overlap=self.config.chunk_overlap
|
||||
)
|
||||
|
||||
# 构建Milvus向量索引
|
||||
print("构建Milvus向量索引...")
|
||||
if not self.index_module.build_vector_index(chunks):
|
||||
raise Exception("构建向量索引失败")
|
||||
|
||||
# 初始化检索器
|
||||
self._initialize_retrievers(chunks)
|
||||
|
||||
# 显示统计信息
|
||||
self._show_knowledge_base_stats()
|
||||
|
||||
print("✅ 知识库构建完成！")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"知识库构建失败: {e}")
|
||||
raise
|
||||
|
||||
def _initialize_retrievers(self, chunks: List = None):
|
||||
"""初始化检索器"""
|
||||
print("初始化检索引擎...")
|
||||
|
||||
# 如果没有chunks，从数据模块获取
|
||||
if chunks is None:
|
||||
chunks = self.data_module.chunks or []
|
||||
|
||||
# 初始化传统检索器
|
||||
self.traditional_retrieval.initialize(chunks)
|
||||
|
||||
# 初始化图RAG检索器
|
||||
self.graph_rag_retrieval.initialize()
|
||||
|
||||
self.system_ready = True
|
||||
print("✅ 检索引擎初始化完成！")
|
||||
|
||||
def _show_knowledge_base_stats(self):
|
||||
"""显示知识库统计信息"""
|
||||
print(f"\n知识库统计:")
|
||||
|
||||
# 数据统计
|
||||
stats = self.data_module.get_statistics()
|
||||
print(f" 菜谱数量: {stats.get('total_recipes', 0)}")
|
||||
print(f" 食材数量: {stats.get('total_ingredients', 0)}")
|
||||
print(f" 烹饪步骤: {stats.get('total_cooking_steps', 0)}")
|
||||
print(f" 文档数量: {stats.get('total_documents', 0)}")
|
||||
print(f" 文本块数: {stats.get('total_chunks', 0)}")
|
||||
|
||||
# Milvus统计
|
||||
milvus_stats = self.index_module.get_collection_stats()
|
||||
print(f" 向量索引: {milvus_stats.get('row_count', 0)} 条记录")
|
||||
|
||||
# 图RAG统计
|
||||
route_stats = self.query_router.get_route_statistics()
|
||||
print(f" 路由统计: 总查询 {route_stats.get('total_queries', 0)} 次")
|
||||
|
||||
if stats.get('categories'):
|
||||
categories = list(stats['categories'].keys())[:10]
|
||||
print(f" 🏷️ 主要分类: {', '.join(categories)}")
|
||||
|
||||
def ask_question_with_routing(self, question: str, stream: bool = False, explain_routing: bool = False):
|
||||
"""
|
||||
智能问答：自动选择最佳检索策略
|
||||
"""
|
||||
if not self.system_ready:
|
||||
raise ValueError("系统未就绪，请先构建知识库")
|
||||
|
||||
print(f"\n❓ 用户问题: {question}")
|
||||
|
||||
# 显示路由决策解释（可选）
|
||||
if explain_routing:
|
||||
explanation = self.query_router.explain_routing_decision(question)
|
||||
print(explanation)
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 1. 智能路由检索
|
||||
print("执行智能查询路由...")
|
||||
relevant_docs, analysis = self.query_router.route_query(question, self.config.top_k)
|
||||
|
||||
# 2. 显示路由信息
|
||||
strategy_icons = {
|
||||
"hybrid_traditional": "🔍",
|
||||
"graph_rag": "🕸️",
|
||||
"combined": "🔄"
|
||||
}
|
||||
strategy_icon = strategy_icons.get(analysis.recommended_strategy.value, "❓")
|
||||
print(f"{strategy_icon} 使用策略: {analysis.recommended_strategy.value}")
|
||||
print(f"📊 复杂度: {analysis.query_complexity:.2f}, 关系密集度: {analysis.relationship_intensity:.2f}")
|
||||
|
||||
# 3. 显示检索结果信息
|
||||
if relevant_docs:
|
||||
doc_info = []
|
||||
for doc in relevant_docs:
|
||||
recipe_name = doc.metadata.get('recipe_name', '未知内容')
|
||||
search_type = doc.metadata.get('search_type', doc.metadata.get('route_strategy', 'unknown'))
|
||||
score = doc.metadata.get('final_score', doc.metadata.get('relevance_score', 0))
|
||||
doc_info.append(f"{recipe_name}({search_type}, {score:.3f})")
|
||||
|
||||
print(f"📋 找到 {len(relevant_docs)} 个相关文档: {', '.join(doc_info[:3])}")
|
||||
if len(doc_info) > 3:
|
||||
print(f" 等 {len(relevant_docs)} 个结果...")
|
||||
else:
|
||||
# 保持返回值签名一致：始终返回 (result, analysis)
|
||||
return "抱歉，没有找到相关的烹饪信息。请尝试其他问题。", analysis
|
||||
|
||||
# 4. 生成回答
|
||||
print("🎯 智能生成回答...")
|
||||
|
||||
if stream:
|
||||
try:
|
||||
for chunk_text in self.generation_module.generate_adaptive_answer_stream(question, relevant_docs):
|
||||
print(chunk_text, end="", flush=True)
|
||||
print("\n")
|
||||
result = "流式输出完成"
|
||||
except Exception as stream_error:
|
||||
logger.error(f"流式输出过程中出现错误: {stream_error}")
|
||||
print(f"\n⚠️ 流式输出中断，切换到标准模式...")
|
||||
# 使用非流式作为后备
|
||||
result = self.generation_module.generate_adaptive_answer(question, relevant_docs)
|
||||
else:
|
||||
result = self.generation_module.generate_adaptive_answer(question, relevant_docs)
|
||||
|
||||
# 5. 性能统计
|
||||
end_time = time.time()
|
||||
print(f"\n⏱️ 问答完成，耗时: {end_time - start_time:.2f}秒")
|
||||
|
||||
return result, analysis
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"问答处理失败: {e}")
|
||||
return f"抱歉，处理问题时出现错误：{str(e)}", None
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
def run_interactive(self):
|
||||
"""运行交互式问答"""
|
||||
if not self.system_ready:
|
||||
print("❌ 系统未就绪，请先构建知识库")
|
||||
return
|
||||
|
||||
print("\n欢迎使用尝尝咸淡RAG烹饪助手！")
|
||||
print("可用功能：")
|
||||
print(" - 'stats' : 查看系统统计")
|
||||
print(" - 'rebuild' : 重建知识库")
|
||||
print(" - 'quit' : 退出系统")
|
||||
print("\n" + "="*50)
|
||||
|
||||
while True:
|
||||
try:
|
||||
user_input = input("\n您的问题: ").strip()
|
||||
|
||||
if not user_input:
|
||||
continue
|
||||
|
||||
if user_input.lower() == 'quit':
|
||||
break
|
||||
elif user_input.lower() == 'stats':
|
||||
self._show_system_stats()
|
||||
continue
|
||||
elif user_input.lower() == 'rebuild':
|
||||
self._rebuild_knowledge_base()
|
||||
continue
|
||||
|
||||
# 普通问答 - 使用默认设置
|
||||
use_stream = True # 默认使用流式输出
|
||||
explain_routing = False # 默认不显示路由决策
|
||||
|
||||
print("\n回答:")
|
||||
|
||||
result, analysis = self.ask_question_with_routing(
|
||||
user_input,
|
||||
stream=use_stream,
|
||||
explain_routing=explain_routing
|
||||
)
|
||||
|
||||
if not use_stream and result:
|
||||
print(f"{result}\n")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"处理问题时出错: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
print("\n👋 感谢使用尝尝咸淡RAG烹饪助手！")
|
||||
self._cleanup()
|
||||
|
||||
def _show_system_stats(self):
|
||||
"""显示系统统计信息"""
|
||||
print("\n系统运行统计")
|
||||
print("=" * 40)
|
||||
|
||||
# 路由统计
|
||||
route_stats = self.query_router.get_route_statistics()
|
||||
total_queries = route_stats.get('total_queries', 0)
|
||||
|
||||
if total_queries > 0:
|
||||
print(f"总查询次数: {total_queries}")
|
||||
print(f"传统检索: {route_stats.get('traditional_count', 0)} ({route_stats.get('traditional_ratio', 0):.1%})")
|
||||
print(f"图RAG检索: {route_stats.get('graph_rag_count', 0)} ({route_stats.get('graph_rag_ratio', 0):.1%})")
|
||||
print(f"组合策略: {route_stats.get('combined_count', 0)} ({route_stats.get('combined_ratio', 0):.1%})")
|
||||
else:
|
||||
print("暂无查询记录")
|
||||
|
||||
# 知识库统计
|
||||
self._show_knowledge_base_stats()
|
||||
|
||||
def _rebuild_knowledge_base(self):
|
||||
"""重建知识库"""
|
||||
print("\n准备重建知识库...")
|
||||
|
||||
# 确认操作
|
||||
confirm = input("⚠️ 这将删除现有的向量数据并重新构建，是否继续？(y/N): ").strip().lower()
|
||||
if confirm != 'y':
|
||||
print("❌ 重建操作已取消")
|
||||
return
|
||||
|
||||
try:
|
||||
print("删除现有的Milvus集合...")
|
||||
if self.index_module.delete_collection():
|
||||
print("✅ 现有集合已删除")
|
||||
else:
|
||||
print("删除集合时出现问题，继续重建...")
|
||||
|
||||
# 重新构建知识库
|
||||
print("开始重建知识库...")
|
||||
self.build_knowledge_base()
|
||||
|
||||
print("✅ 知识库重建完成！")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"重建知识库失败: {e}")
|
||||
print(f"❌ 重建失败: {e}")
|
||||
print("建议：请检查Milvus服务状态后重试")
|
||||
|
||||
def _cleanup(self):
|
||||
"""清理资源"""
|
||||
if self.data_module:
|
||||
self.data_module.close()
|
||||
if self.traditional_retrieval:
|
||||
self.traditional_retrieval.close()
|
||||
if self.graph_rag_retrieval:
|
||||
self.graph_rag_retrieval.close()
|
||||
if self.index_module:
|
||||
self.index_module.close()
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
try:
|
||||
print("启动高级图RAG系统...")
|
||||
|
||||
# 创建高级图RAG系统
|
||||
rag_system = AdvancedGraphRAGSystem()
|
||||
|
||||
# 初始化系统
|
||||
rag_system.initialize_system()
|
||||
|
||||
# 构建知识库
|
||||
rag_system.build_knowledge_base()
|
||||
|
||||
# 运行交互式问答
|
||||
rag_system.run_interactive()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"系统运行失败: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
print(f"\n❌ 系统错误: {e}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,15 @@
|
||||
"""
|
||||
基于图数据库的RAG模块包
|
||||
"""
|
||||
|
||||
from .graph_data_preparation import GraphDataPreparationModule
|
||||
from .milvus_index_construction import MilvusIndexConstructionModule
|
||||
from .hybrid_retrieval import HybridRetrievalModule
|
||||
from .generation_integration import GenerationIntegrationModule
|
||||
|
||||
__all__ = [
|
||||
'GraphDataPreparationModule',
|
||||
'MilvusIndexConstructionModule',
|
||||
'HybridRetrievalModule',
|
||||
'GenerationIntegrationModule'
|
||||
]
|
||||
@@ -0,0 +1,175 @@
|
||||
"""
|
||||
生成集成模块
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import List
|
||||
|
||||
from openai import OpenAI
|
||||
from langchain_core.documents import Document
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class GenerationIntegrationModule:
|
||||
"""生成集成模块 - 负责答案生成"""
|
||||
|
||||
def __init__(self, model_name: str = "kimi-k2-0711-preview", temperature: float = 0.1, max_tokens: int = 2048):
|
||||
"""
|
||||
初始化生成集成模块
|
||||
"""
|
||||
self.model_name = model_name
|
||||
self.temperature = temperature
|
||||
self.max_tokens = max_tokens
|
||||
|
||||
# 初始化OpenAI客户端（使用Moonshot API）
|
||||
api_key = os.getenv("MOONSHOT_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("请设置 MOONSHOT_API_KEY 环境变量")
|
||||
|
||||
self.client = OpenAI(
|
||||
api_key=api_key,
|
||||
base_url="https://api.moonshot.cn/v1"
|
||||
)
|
||||
|
||||
logger.info(f"生成模块初始化完成，模型: {model_name}")
|
||||
|
||||
def generate_adaptive_answer(self, question: str, documents: List[Document]) -> str:
|
||||
"""
|
||||
智能统一答案生成
|
||||
自动适应不同类型的查询，无需预先分类
|
||||
"""
|
||||
# 构建上下文
|
||||
context_parts = []
|
||||
|
||||
for doc in documents:
|
||||
content = doc.page_content.strip()
|
||||
if content:
|
||||
# 添加检索层级信息（如果有的话）
|
||||
level = doc.metadata.get('retrieval_level', '')
|
||||
if level:
|
||||
context_parts.append(f"[{level.upper()}] {content}")
|
||||
else:
|
||||
context_parts.append(content)
|
||||
|
||||
context = "\n\n".join(context_parts)
|
||||
|
||||
# LightRAG风格的统一提示词
|
||||
prompt = f"""
|
||||
作为一位专业的烹饪助手，请基于以下信息回答用户的问题。
|
||||
|
||||
检索到的相关信息：
|
||||
{context}
|
||||
|
||||
用户问题：{question}
|
||||
|
||||
请提供准确、实用的回答。根据问题的性质：
|
||||
- 如果是询问多个菜品，请提供清晰的列表
|
||||
- 如果是询问具体制作方法，请提供详细步骤
|
||||
- 如果是一般性咨询，请提供综合性回答
|
||||
|
||||
回答：
|
||||
"""
|
||||
|
||||
try:
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=[
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=self.temperature,
|
||||
max_tokens=self.max_tokens
|
||||
)
|
||||
|
||||
return response.choices[0].message.content.strip()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"LightRAG答案生成失败: {e}")
|
||||
return f"抱歉，生成回答时出现错误：{str(e)}"
|
||||
|
||||
def generate_adaptive_answer_stream(self, question: str, documents: List[Document], max_retries: int = 3):
|
||||
"""
|
||||
LightRAG风格的流式答案生成（带重试机制）
|
||||
"""
|
||||
# 构建上下文
|
||||
context_parts = []
|
||||
|
||||
for doc in documents:
|
||||
content = doc.page_content.strip()
|
||||
if content:
|
||||
level = doc.metadata.get('retrieval_level', '')
|
||||
if level:
|
||||
context_parts.append(f"[{level.upper()}] {content}")
|
||||
else:
|
||||
context_parts.append(content)
|
||||
|
||||
context = "\n\n".join(context_parts)
|
||||
|
||||
# LightRAG风格的统一提示词
|
||||
prompt = f"""
|
||||
作为一位专业的烹饪助手，请基于以下信息回答用户的问题。
|
||||
|
||||
检索到的相关信息：
|
||||
{context}
|
||||
|
||||
用户问题：{question}
|
||||
|
||||
请提供准确、实用的回答。根据问题的性质：
|
||||
- 如果是询问多个菜品，请提供清晰的列表
|
||||
- 如果是询问具体制作方法，请提供详细步骤
|
||||
- 如果是一般性咨询，请提供综合性回答
|
||||
|
||||
回答：
|
||||
"""
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=[
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=self.temperature,
|
||||
max_tokens=self.max_tokens,
|
||||
stream=True,
|
||||
timeout=60 # 增加超时设置
|
||||
)
|
||||
|
||||
if attempt == 0:
|
||||
print("开始流式生成回答...\n")
|
||||
else:
|
||||
print(f"第{attempt + 1}次尝试流式生成...\n")
|
||||
|
||||
full_response = ""
|
||||
for chunk in response:
|
||||
if chunk.choices[0].delta.content:
|
||||
content = chunk.choices[0].delta.content
|
||||
full_response += content
|
||||
yield content # 使用yield返回流式内容
|
||||
|
||||
# 如果成功完成，退出重试循环
|
||||
return
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"流式生成第{attempt + 1}次尝试失败: {e}")
|
||||
|
||||
if attempt < max_retries - 1:
|
||||
wait_time = (attempt + 1) * 2 # 递增等待时间
|
||||
print(f"⚠️ 连接中断，{wait_time}秒后重试...")
|
||||
time.sleep(wait_time)
|
||||
continue
|
||||
else:
|
||||
# 所有重试都失败，使用非流式作为后备
|
||||
logger.error(f"流式生成完全失败，尝试非流式后备方案")
|
||||
print("⚠️ 流式生成失败，切换到标准模式...")
|
||||
|
||||
try:
|
||||
fallback_response = self.generate_adaptive_answer(question, documents)
|
||||
yield fallback_response
|
||||
return
|
||||
except Exception as fallback_error:
|
||||
logger.error(f"后备生成也失败: {fallback_error}")
|
||||
error_msg = f"抱歉，生成回答时出现网络错误，请稍后重试。错误信息：{str(e)}"
|
||||
yield error_msg
|
||||
return
|
||||
@@ -0,0 +1,458 @@
|
||||
"""
|
||||
图数据库数据准备模块
|
||||
"""
|
||||
|
||||
import logging
|
||||
import json
|
||||
from typing import List, Dict, Any, Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
from neo4j import GraphDatabase
|
||||
from langchain_core.documents import Document
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@dataclass
|
||||
class GraphNode:
|
||||
"""图节点数据结构"""
|
||||
node_id: str
|
||||
labels: List[str]
|
||||
name: str
|
||||
properties: Dict[str, Any]
|
||||
|
||||
@dataclass
|
||||
class GraphRelation:
|
||||
"""图关系数据结构"""
|
||||
start_node_id: str
|
||||
end_node_id: str
|
||||
relation_type: str
|
||||
properties: Dict[str, Any]
|
||||
|
||||
class GraphDataPreparationModule:
|
||||
"""图数据库数据准备模块 - 从Neo4j读取数据并转换为文档"""
|
||||
|
||||
def __init__(self, uri: str, user: str, password: str, database: str = "neo4j"):
|
||||
"""
|
||||
初始化图数据库连接
|
||||
|
||||
Args:
|
||||
uri: Neo4j连接URI
|
||||
user: 用户名
|
||||
password: 密码
|
||||
database: 数据库名称
|
||||
"""
|
||||
self.uri = uri
|
||||
self.user = user
|
||||
self.password = password
|
||||
self.database = database
|
||||
self.driver = None
|
||||
self.documents: List[Document] = []
|
||||
self.chunks: List[Document] = []
|
||||
self.recipes: List[GraphNode] = []
|
||||
self.ingredients: List[GraphNode] = []
|
||||
self.cooking_steps: List[GraphNode] = []
|
||||
|
||||
self._connect()
|
||||
|
||||
def _connect(self):
|
||||
"""建立Neo4j连接"""
|
||||
try:
|
||||
self.driver = GraphDatabase.driver(
|
||||
self.uri,
|
||||
auth=(self.user, self.password),
|
||||
database=self.database
|
||||
)
|
||||
logger.info(f"已连接到Neo4j数据库: {self.uri}")
|
||||
|
||||
# 测试连接
|
||||
with self.driver.session() as session:
|
||||
result = session.run("RETURN 1 as test")
|
||||
test_result = result.single()
|
||||
if test_result:
|
||||
logger.info("Neo4j连接测试成功")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"连接Neo4j失败: {e}")
|
||||
raise
|
||||
|
||||
def close(self):
|
||||
"""关闭数据库连接"""
|
||||
if hasattr(self, 'driver') and self.driver:
|
||||
self.driver.close()
|
||||
logger.info("Neo4j连接已关闭")
|
||||
|
||||
def load_graph_data(self) -> Dict[str, Any]:
|
||||
"""
|
||||
从Neo4j加载图数据
|
||||
|
||||
Returns:
|
||||
包含节点和关系的数据字典
|
||||
"""
|
||||
logger.info("正在从Neo4j加载图数据...")
|
||||
|
||||
with self.driver.session() as session:
|
||||
# 加载所有菜谱节点，从Category关系中读取分类信息
|
||||
recipes_query = """
|
||||
MATCH (r:Recipe)
|
||||
WHERE r.nodeId >= '200000000'
|
||||
OPTIONAL MATCH (r)-[:BELONGS_TO_CATEGORY]->(c:Category)
|
||||
WITH r, collect(c.name) as categories
|
||||
RETURN r.nodeId as nodeId, labels(r) as labels, r.name as name,
|
||||
properties(r) as originalProperties,
|
||||
CASE WHEN size(categories) > 0
|
||||
THEN categories[0]
|
||||
ELSE COALESCE(r.category, '未知') END as mainCategory,
|
||||
CASE WHEN size(categories) > 0
|
||||
THEN categories
|
||||
ELSE [COALESCE(r.category, '未知')] END as allCategories
|
||||
ORDER BY r.nodeId
|
||||
"""
|
||||
|
||||
result = session.run(recipes_query)
|
||||
self.recipes = []
|
||||
for record in result:
|
||||
# 合并原始属性和新的分类信息
|
||||
properties = dict(record["originalProperties"])
|
||||
properties["category"] = record["mainCategory"]
|
||||
properties["all_categories"] = record["allCategories"]
|
||||
|
||||
node = GraphNode(
|
||||
node_id=record["nodeId"],
|
||||
labels=record["labels"],
|
||||
name=record["name"],
|
||||
properties=properties
|
||||
)
|
||||
self.recipes.append(node)
|
||||
|
||||
logger.info(f"加载了 {len(self.recipes)} 个菜谱节点")
|
||||
|
||||
# 加载所有食材节点
|
||||
ingredients_query = """
|
||||
MATCH (i:Ingredient)
|
||||
WHERE i.nodeId >= '200000000'
|
||||
RETURN i.nodeId as nodeId, labels(i) as labels, i.name as name,
|
||||
properties(i) as properties
|
||||
ORDER BY i.nodeId
|
||||
"""
|
||||
|
||||
result = session.run(ingredients_query)
|
||||
self.ingredients = []
|
||||
for record in result:
|
||||
node = GraphNode(
|
||||
node_id=record["nodeId"],
|
||||
labels=record["labels"],
|
||||
name=record["name"],
|
||||
properties=record["properties"]
|
||||
)
|
||||
self.ingredients.append(node)
|
||||
|
||||
logger.info(f"加载了 {len(self.ingredients)} 个食材节点")
|
||||
|
||||
# 加载所有烹饪步骤节点
|
||||
steps_query = """
|
||||
MATCH (s:CookingStep)
|
||||
WHERE s.nodeId >= '200000000'
|
||||
RETURN s.nodeId as nodeId, labels(s) as labels, s.name as name,
|
||||
properties(s) as properties
|
||||
ORDER BY s.nodeId
|
||||
"""
|
||||
|
||||
result = session.run(steps_query)
|
||||
self.cooking_steps = []
|
||||
for record in result:
|
||||
node = GraphNode(
|
||||
node_id=record["nodeId"],
|
||||
labels=record["labels"],
|
||||
name=record["name"],
|
||||
properties=record["properties"]
|
||||
)
|
||||
self.cooking_steps.append(node)
|
||||
|
||||
logger.info(f"加载了 {len(self.cooking_steps)} 个烹饪步骤节点")
|
||||
|
||||
return {
|
||||
'recipes': len(self.recipes),
|
||||
'ingredients': len(self.ingredients),
|
||||
'cooking_steps': len(self.cooking_steps)
|
||||
}
|
||||
|
||||
def build_recipe_documents(self) -> List[Document]:
|
||||
"""
|
||||
构建菜谱文档，集成相关的食材和步骤信息
|
||||
|
||||
Returns:
|
||||
结构化的菜谱文档列表
|
||||
"""
|
||||
logger.info("正在构建菜谱文档...")
|
||||
|
||||
documents = []
|
||||
|
||||
with self.driver.session() as session:
|
||||
for recipe in self.recipes:
|
||||
try:
|
||||
recipe_id = recipe.node_id
|
||||
recipe_name = recipe.name
|
||||
|
||||
# 获取菜谱的相关食材
|
||||
ingredients_query = """
|
||||
MATCH (r:Recipe {nodeId: $recipe_id})-[req:REQUIRES]->(i:Ingredient)
|
||||
RETURN i.name as name, i.category as category,
|
||||
req.amount as amount, req.unit as unit,
|
||||
i.description as description
|
||||
ORDER BY i.name
|
||||
"""
|
||||
|
||||
ingredients_result = session.run(ingredients_query, {"recipe_id": recipe_id})
|
||||
ingredients_info = []
|
||||
for ing_record in ingredients_result:
|
||||
amount = ing_record.get("amount", "")
|
||||
unit = ing_record.get("unit", "")
|
||||
ingredient_text = f"{ing_record['name']}"
|
||||
if amount and unit:
|
||||
ingredient_text += f"({amount}{unit})"
|
||||
if ing_record.get("description"):
|
||||
ingredient_text += f" - {ing_record['description']}"
|
||||
ingredients_info.append(ingredient_text)
|
||||
|
||||
# 获取菜谱的烹饪步骤
|
||||
steps_query = """
|
||||
MATCH (r:Recipe {nodeId: $recipe_id})-[c:CONTAINS_STEP]->(s:CookingStep)
|
||||
RETURN s.name as name, s.description as description,
|
||||
s.stepNumber as stepNumber, s.methods as methods,
|
||||
s.tools as tools, s.timeEstimate as timeEstimate,
|
||||
c.stepOrder as stepOrder
|
||||
ORDER BY COALESCE(c.stepOrder, s.stepNumber, 999)
|
||||
"""
|
||||
|
||||
steps_result = session.run(steps_query, {"recipe_id": recipe_id})
|
||||
steps_info = []
|
||||
for step_record in steps_result:
|
||||
step_text = f"步骤: {step_record['name']}"
|
||||
if step_record.get("description"):
|
||||
step_text += f"\n描述: {step_record['description']}"
|
||||
if step_record.get("methods"):
|
||||
step_text += f"\n方法: {step_record['methods']}"
|
||||
if step_record.get("tools"):
|
||||
step_text += f"\n工具: {step_record['tools']}"
|
||||
if step_record.get("timeEstimate"):
|
||||
step_text += f"\n时间: {step_record['timeEstimate']}"
|
||||
steps_info.append(step_text)
|
||||
|
||||
# 构建完整的菜谱文档内容
|
||||
content_parts = [f"# {recipe_name}"]
|
||||
|
||||
# 添加菜谱基本信息
|
||||
if recipe.properties.get("description"):
|
||||
content_parts.append(f"\n## 菜品描述\n{recipe.properties['description']}")
|
||||
|
||||
if recipe.properties.get("cuisineType"):
|
||||
content_parts.append(f"\n菜系: {recipe.properties['cuisineType']}")
|
||||
|
||||
if recipe.properties.get("difficulty"):
|
||||
content_parts.append(f"难度: {recipe.properties['difficulty']}星")
|
||||
|
||||
if recipe.properties.get("prepTime") or recipe.properties.get("cookTime"):
|
||||
time_info = []
|
||||
if recipe.properties.get("prepTime"):
|
||||
time_info.append(f"准备时间: {recipe.properties['prepTime']}")
|
||||
if recipe.properties.get("cookTime"):
|
||||
time_info.append(f"烹饪时间: {recipe.properties['cookTime']}")
|
||||
content_parts.append(f"\n时间信息: {', '.join(time_info)}")
|
||||
|
||||
if recipe.properties.get("servings"):
|
||||
content_parts.append(f"份量: {recipe.properties['servings']}")
|
||||
|
||||
# 添加食材信息
|
||||
if ingredients_info:
|
||||
content_parts.append("\n## 所需食材")
|
||||
for i, ingredient in enumerate(ingredients_info, 1):
|
||||
content_parts.append(f"{i}. {ingredient}")
|
||||
|
||||
# 添加步骤信息
|
||||
if steps_info:
|
||||
content_parts.append("\n## 制作步骤")
|
||||
for i, step in enumerate(steps_info, 1):
|
||||
content_parts.append(f"\n### 第{i}步\n{step}")
|
||||
|
||||
# 添加标签信息
|
||||
if recipe.properties.get("tags"):
|
||||
content_parts.append(f"\n## 标签\n{recipe.properties['tags']}")
|
||||
|
||||
# 组合成最终内容
|
||||
full_content = "\n".join(content_parts)
|
||||
|
||||
# 创建文档对象
|
||||
doc = Document(
|
||||
page_content=full_content,
|
||||
metadata={
|
||||
"node_id": recipe_id,
|
||||
"recipe_name": recipe_name,
|
||||
"node_type": "Recipe",
|
||||
"category": recipe.properties.get("category", "未知"),
|
||||
"cuisine_type": recipe.properties.get("cuisineType", "未知"),
|
||||
"difficulty": recipe.properties.get("difficulty", 0),
|
||||
"prep_time": recipe.properties.get("prepTime", ""),
|
||||
"cook_time": recipe.properties.get("cookTime", ""),
|
||||
"servings": recipe.properties.get("servings", ""),
|
||||
"ingredients_count": len(ingredients_info),
|
||||
"steps_count": len(steps_info),
|
||||
"doc_type": "recipe",
|
||||
"content_length": len(full_content)
|
||||
}
|
||||
)
|
||||
|
||||
documents.append(doc)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"构建菜谱文档失败 {recipe_name} (ID: {recipe_id}): {e}")
|
||||
continue
|
||||
|
||||
self.documents = documents
|
||||
logger.info(f"成功构建 {len(documents)} 个菜谱文档")
|
||||
return documents
|
||||
|
||||
def chunk_documents(self, chunk_size: int = 500, chunk_overlap: int = 50) -> List[Document]:
|
||||
"""
|
||||
对文档进行分块处理
|
||||
|
||||
Args:
|
||||
chunk_size: 分块大小
|
||||
chunk_overlap: 重叠大小
|
||||
|
||||
Returns:
|
||||
分块后的文档列表
|
||||
"""
|
||||
logger.info(f"正在进行文档分块，块大小: {chunk_size}, 重叠: {chunk_overlap}")
|
||||
|
||||
if not self.documents:
|
||||
raise ValueError("请先构建文档")
|
||||
|
||||
chunks = []
|
||||
chunk_id = 0
|
||||
|
||||
for doc in self.documents:
|
||||
content = doc.page_content
|
||||
|
||||
# 简单的按长度分块
|
||||
if len(content) <= chunk_size:
|
||||
# 内容较短，不需要分块
|
||||
chunk = Document(
|
||||
page_content=content,
|
||||
metadata={
|
||||
**doc.metadata,
|
||||
"chunk_id": f"{doc.metadata['node_id']}_chunk_{chunk_id}",
|
||||
"parent_id": doc.metadata["node_id"],
|
||||
"chunk_index": 0,
|
||||
"total_chunks": 1,
|
||||
"chunk_size": len(content),
|
||||
"doc_type": "chunk"
|
||||
}
|
||||
)
|
||||
chunks.append(chunk)
|
||||
chunk_id += 1
|
||||
else:
|
||||
# 按章节分块（基于标题）
|
||||
sections = content.split('\n## ')
|
||||
if len(sections) <= 1:
|
||||
# 没有二级标题，按长度强制分块
|
||||
total_chunks = (len(content) - 1) // (chunk_size - chunk_overlap) + 1
|
||||
|
||||
for i in range(total_chunks):
|
||||
start = i * (chunk_size - chunk_overlap)
|
||||
end = min(start + chunk_size, len(content))
|
||||
|
||||
chunk_content = content[start:end]
|
||||
|
||||
chunk = Document(
|
||||
page_content=chunk_content,
|
||||
metadata={
|
||||
**doc.metadata,
|
||||
"chunk_id": f"{doc.metadata['node_id']}_chunk_{chunk_id}",
|
||||
"parent_id": doc.metadata["node_id"],
|
||||
"chunk_index": i,
|
||||
"total_chunks": total_chunks,
|
||||
"chunk_size": len(chunk_content),
|
||||
"doc_type": "chunk"
|
||||
}
|
||||
)
|
||||
chunks.append(chunk)
|
||||
chunk_id += 1
|
||||
else:
|
||||
# 按章节分块
|
||||
total_chunks = len(sections)
|
||||
for i, section in enumerate(sections):
|
||||
if i == 0:
|
||||
# 第一个部分包含标题
|
||||
chunk_content = section
|
||||
else:
|
||||
# 其他部分添加章节标题
|
||||
chunk_content = f"## {section}"
|
||||
|
||||
chunk = Document(
|
||||
page_content=chunk_content,
|
||||
metadata={
|
||||
**doc.metadata,
|
||||
"chunk_id": f"{doc.metadata['node_id']}_chunk_{chunk_id}",
|
||||
"parent_id": doc.metadata["node_id"],
|
||||
"chunk_index": i,
|
||||
"total_chunks": total_chunks,
|
||||
"chunk_size": len(chunk_content),
|
||||
"doc_type": "chunk",
|
||||
"section_title": section.split('\n')[0] if i > 0 else "主标题"
|
||||
}
|
||||
)
|
||||
chunks.append(chunk)
|
||||
chunk_id += 1
|
||||
|
||||
self.chunks = chunks
|
||||
logger.info(f"文档分块完成，共生成 {len(chunks)} 个块")
|
||||
return chunks
|
||||
|
||||
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""
|
||||
获取数据统计信息
|
||||
|
||||
Returns:
|
||||
统计信息字典
|
||||
"""
|
||||
stats = {
|
||||
'total_recipes': len(self.recipes),
|
||||
'total_ingredients': len(self.ingredients),
|
||||
'total_cooking_steps': len(self.cooking_steps),
|
||||
'total_documents': len(self.documents),
|
||||
'total_chunks': len(self.chunks)
|
||||
}
|
||||
|
||||
if self.documents:
|
||||
# 分类统计
|
||||
categories = {}
|
||||
cuisines = {}
|
||||
difficulties = {}
|
||||
|
||||
for doc in self.documents:
|
||||
category = doc.metadata.get('category', '未知')
|
||||
categories[category] = categories.get(category, 0) + 1
|
||||
|
||||
cuisine = doc.metadata.get('cuisine_type', '未知')
|
||||
cuisines[cuisine] = cuisines.get(cuisine, 0) + 1
|
||||
|
||||
difficulty = doc.metadata.get('difficulty', 0)
|
||||
difficulties[str(difficulty)] = difficulties.get(str(difficulty), 0) + 1
|
||||
|
||||
stats.update({
|
||||
'categories': categories,
|
||||
'cuisines': cuisines,
|
||||
'difficulties': difficulties,
|
||||
'avg_content_length': sum(doc.metadata.get('content_length', 0) for doc in self.documents) / len(self.documents),
|
||||
'avg_chunk_size': sum(chunk.metadata.get('chunk_size', 0) for chunk in self.chunks) / len(self.chunks) if self.chunks else 0
|
||||
})
|
||||
|
||||
return stats
|
||||
|
||||
|
||||
|
||||
def __del__(self):
|
||||
"""析构函数，确保关闭连接"""
|
||||
self.close()
|
||||
@@ -0,0 +1,387 @@
|
||||
"""
|
||||
图索引模块
|
||||
实现实体和关系的键值对结构 (K,V)
|
||||
K: 索引键（简短词汇或短语）
|
||||
V: 详细描述段落（包含相关文本片段）
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Dict, List, Tuple, Any, Optional
|
||||
from dataclasses import dataclass
|
||||
from collections import defaultdict
|
||||
|
||||
from langchain_core.documents import Document
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@dataclass
|
||||
class EntityKeyValue:
|
||||
"""实体键值对"""
|
||||
entity_name: str
|
||||
index_keys: List[str] # 索引键列表
|
||||
value_content: str # 详细描述内容
|
||||
entity_type: str # 实体类型 (Recipe, Ingredient, CookingStep)
|
||||
metadata: Dict[str, Any]
|
||||
|
||||
@dataclass
|
||||
class RelationKeyValue:
|
||||
"""关系键值对"""
|
||||
relation_id: str
|
||||
index_keys: List[str] # 多个索引键（可包含全局主题）
|
||||
value_content: str # 关系描述内容
|
||||
relation_type: str # 关系类型
|
||||
source_entity: str # 源实体
|
||||
target_entity: str # 目标实体
|
||||
metadata: Dict[str, Any]
|
||||
|
||||
class GraphIndexingModule:
|
||||
"""
|
||||
图索引模块
|
||||
核心功能：
|
||||
1. 为实体创建键值对（名称作为唯一索引键）
|
||||
2. 为关系创建键值对（多个索引键，包含全局主题）
|
||||
3. 去重和优化图操作
|
||||
4. 支持增量更新
|
||||
"""
|
||||
|
||||
def __init__(self, config, llm_client):
|
||||
self.config = config
|
||||
self.llm_client = llm_client
|
||||
|
||||
# 键值对存储
|
||||
self.entity_kv_store: Dict[str, EntityKeyValue] = {}
|
||||
self.relation_kv_store: Dict[str, RelationKeyValue] = {}
|
||||
|
||||
# 索引映射：key -> entity/relation IDs
|
||||
self.key_to_entities: Dict[str, List[str]] = defaultdict(list)
|
||||
self.key_to_relations: Dict[str, List[str]] = defaultdict(list)
|
||||
|
||||
def create_entity_key_values(self, recipes: List[Any], ingredients: List[Any],
|
||||
cooking_steps: List[Any]) -> Dict[str, EntityKeyValue]:
|
||||
"""
|
||||
为实体创建键值对结构
|
||||
每个实体使用其名称作为唯一索引键
|
||||
"""
|
||||
logger.info("开始创建实体键值对...")
|
||||
|
||||
# 处理菜谱实体
|
||||
for recipe in recipes:
|
||||
entity_id = recipe.node_id
|
||||
entity_name = recipe.name or f"菜谱_{entity_id}"
|
||||
|
||||
# 构建详细内容
|
||||
content_parts = [f"菜品名称: {entity_name}"]
|
||||
|
||||
if hasattr(recipe, 'properties'):
|
||||
props = recipe.properties
|
||||
if props.get('description'):
|
||||
content_parts.append(f"描述: {props['description']}")
|
||||
if props.get('category'):
|
||||
content_parts.append(f"分类: {props['category']}")
|
||||
if props.get('cuisineType'):
|
||||
content_parts.append(f"菜系: {props['cuisineType']}")
|
||||
if props.get('difficulty'):
|
||||
content_parts.append(f"难度: {props['difficulty']}")
|
||||
if props.get('cookingTime'):
|
||||
content_parts.append(f"制作时间: {props['cookingTime']}")
|
||||
|
||||
# 创建键值对
|
||||
entity_kv = EntityKeyValue(
|
||||
entity_name=entity_name,
|
||||
index_keys=[entity_name], # 使用名称作为唯一索引键
|
||||
value_content='\n'.join(content_parts),
|
||||
entity_type="Recipe",
|
||||
metadata={
|
||||
"node_id": entity_id,
|
||||
"properties": getattr(recipe, 'properties', {})
|
||||
}
|
||||
)
|
||||
|
||||
self.entity_kv_store[entity_id] = entity_kv
|
||||
self.key_to_entities[entity_name].append(entity_id)
|
||||
|
||||
# 处理食材实体
|
||||
for ingredient in ingredients:
|
||||
entity_id = ingredient.node_id
|
||||
entity_name = ingredient.name or f"食材_{entity_id}"
|
||||
|
||||
content_parts = [f"食材名称: {entity_name}"]
|
||||
|
||||
if hasattr(ingredient, 'properties'):
|
||||
props = ingredient.properties
|
||||
if props.get('category'):
|
||||
content_parts.append(f"类别: {props['category']}")
|
||||
if props.get('nutrition'):
|
||||
content_parts.append(f"营养信息: {props['nutrition']}")
|
||||
if props.get('storage'):
|
||||
content_parts.append(f"储存方式: {props['storage']}")
|
||||
|
||||
entity_kv = EntityKeyValue(
|
||||
entity_name=entity_name,
|
||||
index_keys=[entity_name],
|
||||
value_content='\n'.join(content_parts),
|
||||
entity_type="Ingredient",
|
||||
metadata={
|
||||
"node_id": entity_id,
|
||||
"properties": getattr(ingredient, 'properties', {})
|
||||
}
|
||||
)
|
||||
|
||||
self.entity_kv_store[entity_id] = entity_kv
|
||||
self.key_to_entities[entity_name].append(entity_id)
|
||||
|
||||
# 处理烹饪步骤实体
|
||||
for step in cooking_steps:
|
||||
entity_id = step.node_id
|
||||
entity_name = f"步骤_{entity_id}"
|
||||
|
||||
content_parts = [f"烹饪步骤: {entity_name}"]
|
||||
|
||||
if hasattr(step, 'properties'):
|
||||
props = step.properties
|
||||
if props.get('description'):
|
||||
content_parts.append(f"步骤描述: {props['description']}")
|
||||
if props.get('order'):
|
||||
content_parts.append(f"步骤顺序: {props['order']}")
|
||||
if props.get('technique'):
|
||||
content_parts.append(f"技巧: {props['technique']}")
|
||||
if props.get('time'):
|
||||
content_parts.append(f"时间: {props['time']}")
|
||||
|
||||
entity_kv = EntityKeyValue(
|
||||
entity_name=entity_name,
|
||||
index_keys=[entity_name],
|
||||
value_content='\n'.join(content_parts),
|
||||
entity_type="CookingStep",
|
||||
metadata={
|
||||
"node_id": entity_id,
|
||||
"properties": getattr(step, 'properties', {})
|
||||
}
|
||||
)
|
||||
|
||||
self.entity_kv_store[entity_id] = entity_kv
|
||||
self.key_to_entities[entity_name].append(entity_id)
|
||||
|
||||
logger.info(f"实体键值对创建完成，共 {len(self.entity_kv_store)} 个实体")
|
||||
return self.entity_kv_store
|
||||
|
||||
def create_relation_key_values(self, relationships: List[Tuple[str, str, str]]) -> Dict[str, RelationKeyValue]:
|
||||
"""
|
||||
为关系创建键值对结构
|
||||
关系可能有多个索引键，包含从LLM增强的全局主题
|
||||
"""
|
||||
logger.info("开始创建关系键值对...")
|
||||
|
||||
for i, (source_id, relation_type, target_id) in enumerate(relationships):
|
||||
relation_id = f"rel_{i}_{source_id}_{target_id}"
|
||||
|
||||
# 获取源实体和目标实体信息
|
||||
source_entity = self.entity_kv_store.get(source_id)
|
||||
target_entity = self.entity_kv_store.get(target_id)
|
||||
|
||||
if not source_entity or not target_entity:
|
||||
continue
|
||||
|
||||
# 构建关系描述
|
||||
content_parts = [
|
||||
f"关系类型: {relation_type}",
|
||||
f"源实体: {source_entity.entity_name} ({source_entity.entity_type})",
|
||||
f"目标实体: {target_entity.entity_name} ({target_entity.entity_type})"
|
||||
]
|
||||
|
||||
# 生成多个索引键（包含全局主题）
|
||||
index_keys = self._generate_relation_index_keys(
|
||||
source_entity, target_entity, relation_type
|
||||
)
|
||||
|
||||
# 创建关系键值对
|
||||
relation_kv = RelationKeyValue(
|
||||
relation_id=relation_id,
|
||||
index_keys=index_keys,
|
||||
value_content='\n'.join(content_parts),
|
||||
relation_type=relation_type,
|
||||
source_entity=source_id,
|
||||
target_entity=target_id,
|
||||
metadata={
|
||||
"source_name": source_entity.entity_name,
|
||||
"target_name": target_entity.entity_name,
|
||||
"created_from_graph": True
|
||||
}
|
||||
)
|
||||
|
||||
self.relation_kv_store[relation_id] = relation_kv
|
||||
|
||||
# 为每个索引键建立映射
|
||||
for key in index_keys:
|
||||
self.key_to_relations[key].append(relation_id)
|
||||
|
||||
logger.info(f"关系键值对创建完成，共 {len(self.relation_kv_store)} 个关系")
|
||||
return self.relation_kv_store
|
||||
|
||||
def _generate_relation_index_keys(self, source_entity: EntityKeyValue,
|
||||
target_entity: EntityKeyValue,
|
||||
relation_type: str) -> List[str]:
|
||||
"""
|
||||
为关系生成多个索引键，包含全局主题
|
||||
"""
|
||||
keys = [relation_type] # 基础关系类型键
|
||||
|
||||
# 根据关系类型和实体类型生成主题键
|
||||
if relation_type == "REQUIRES":
|
||||
# 菜谱-食材关系的主题键
|
||||
keys.extend([
|
||||
"食材搭配",
|
||||
"烹饪原料",
|
||||
f"{source_entity.entity_name}_食材",
|
||||
target_entity.entity_name
|
||||
])
|
||||
elif relation_type == "HAS_STEP":
|
||||
# 菜谱-步骤关系的主题键
|
||||
keys.extend([
|
||||
"制作步骤",
|
||||
"烹饪过程",
|
||||
f"{source_entity.entity_name}_步骤",
|
||||
"制作方法"
|
||||
])
|
||||
elif relation_type == "BELONGS_TO_CATEGORY":
|
||||
# 分类关系的主题键
|
||||
keys.extend([
|
||||
"菜品分类",
|
||||
"美食类别",
|
||||
target_entity.entity_name
|
||||
])
|
||||
|
||||
# 使用LLM增强关系索引键（可选）
|
||||
if getattr(self.config, 'enable_llm_relation_keys', False):
|
||||
enhanced_keys = self._llm_enhance_relation_keys(source_entity, target_entity, relation_type)
|
||||
keys.extend(enhanced_keys)
|
||||
|
||||
# 去重并返回
|
||||
return list(set(keys))
|
||||
|
||||
def _llm_enhance_relation_keys(self, source_entity: EntityKeyValue,
|
||||
target_entity: EntityKeyValue,
|
||||
relation_type: str) -> List[str]:
|
||||
"""
|
||||
使用LLM增强关系索引键，生成全局主题
|
||||
"""
|
||||
prompt = f"""
|
||||
分析以下实体关系，生成相关的主题关键词：
|
||||
|
||||
源实体: {source_entity.entity_name} ({source_entity.entity_type})
|
||||
目标实体: {target_entity.entity_name} ({target_entity.entity_type})
|
||||
关系类型: {relation_type}
|
||||
|
||||
请生成3-5个相关的主题关键词，用于索引和检索。
|
||||
返回JSON格式：{{"keywords": ["关键词1", "关键词2", "关键词3"]}}
|
||||
"""
|
||||
|
||||
try:
|
||||
response = self.llm_client.chat.completions.create(
|
||||
model=self.config.llm_model,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
temperature=0.1,
|
||||
max_tokens=200
|
||||
)
|
||||
|
||||
result = json.loads(response.choices[0].message.content.strip())
|
||||
return result.get("keywords", [])
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"LLM增强关系索引键失败: {e}")
|
||||
return []
|
||||
|
||||
def deduplicate_entities_and_relations(self):
|
||||
"""
|
||||
去重相同的实体和关系，优化图操作
|
||||
"""
|
||||
logger.info("开始去重实体和关系...")
|
||||
|
||||
# 实体去重：基于名称
|
||||
name_to_entities = defaultdict(list)
|
||||
for entity_id, entity_kv in self.entity_kv_store.items():
|
||||
name_to_entities[entity_kv.entity_name].append(entity_id)
|
||||
|
||||
# 合并重复实体
|
||||
entities_to_remove = []
|
||||
for name, entity_ids in name_to_entities.items():
|
||||
if len(entity_ids) > 1:
|
||||
# 保留第一个，合并其他的内容
|
||||
primary_id = entity_ids[0]
|
||||
primary_entity = self.entity_kv_store[primary_id]
|
||||
|
||||
for entity_id in entity_ids[1:]:
|
||||
duplicate_entity = self.entity_kv_store[entity_id]
|
||||
# 合并内容
|
||||
primary_entity.value_content += f"\n\n补充信息: {duplicate_entity.value_content}"
|
||||
# 标记删除
|
||||
entities_to_remove.append(entity_id)
|
||||
|
||||
# 删除重复实体
|
||||
for entity_id in entities_to_remove:
|
||||
del self.entity_kv_store[entity_id]
|
||||
|
||||
# 关系去重：基于源-目标-类型
|
||||
relation_signature_to_ids = defaultdict(list)
|
||||
for relation_id, relation_kv in self.relation_kv_store.items():
|
||||
signature = f"{relation_kv.source_entity}_{relation_kv.target_entity}_{relation_kv.relation_type}"
|
||||
relation_signature_to_ids[signature].append(relation_id)
|
||||
|
||||
# 合并重复关系
|
||||
relations_to_remove = []
|
||||
for signature, relation_ids in relation_signature_to_ids.items():
|
||||
if len(relation_ids) > 1:
|
||||
# 保留第一个，删除其他
|
||||
for relation_id in relation_ids[1:]:
|
||||
relations_to_remove.append(relation_id)
|
||||
|
||||
# 删除重复关系
|
||||
for relation_id in relations_to_remove:
|
||||
del self.relation_kv_store[relation_id]
|
||||
|
||||
# 重建索引映射
|
||||
self._rebuild_key_mappings()
|
||||
|
||||
logger.info(f"去重完成 - 删除了 {len(entities_to_remove)} 个重复实体，{len(relations_to_remove)} 个重复关系")
|
||||
|
||||
def _rebuild_key_mappings(self):
|
||||
"""重建键到实体/关系的映射"""
|
||||
self.key_to_entities.clear()
|
||||
self.key_to_relations.clear()
|
||||
|
||||
# 重建实体映射
|
||||
for entity_id, entity_kv in self.entity_kv_store.items():
|
||||
for key in entity_kv.index_keys:
|
||||
self.key_to_entities[key].append(entity_id)
|
||||
|
||||
# 重建关系映射
|
||||
for relation_id, relation_kv in self.relation_kv_store.items():
|
||||
for key in relation_kv.index_keys:
|
||||
self.key_to_relations[key].append(relation_id)
|
||||
|
||||
def get_entities_by_key(self, key: str) -> List[EntityKeyValue]:
|
||||
"""根据索引键获取实体"""
|
||||
entity_ids = self.key_to_entities.get(key, [])
|
||||
return [self.entity_kv_store[eid] for eid in entity_ids if eid in self.entity_kv_store]
|
||||
|
||||
def get_relations_by_key(self, key: str) -> List[RelationKeyValue]:
|
||||
"""根据索引键获取关系"""
|
||||
relation_ids = self.key_to_relations.get(key, [])
|
||||
return [self.relation_kv_store[rid] for rid in relation_ids if rid in self.relation_kv_store]
|
||||
|
||||
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""获取键值对存储统计信息"""
|
||||
return {
|
||||
"total_entities": len(self.entity_kv_store),
|
||||
"total_relations": len(self.relation_kv_store),
|
||||
"total_entity_keys": sum(len(kv.index_keys) for kv in self.entity_kv_store.values()),
|
||||
"total_relation_keys": sum(len(kv.index_keys) for kv in self.relation_kv_store.values()),
|
||||
"entity_types": {
|
||||
"Recipe": len([kv for kv in self.entity_kv_store.values() if kv.entity_type == "Recipe"]),
|
||||
"Ingredient": len([kv for kv in self.entity_kv_store.values() if kv.entity_type == "Ingredient"]),
|
||||
"CookingStep": len([kv for kv in self.entity_kv_store.values() if kv.entity_type == "CookingStep"])
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user