UPDATE: RAG process

retrieval_augmented_generation/build_embeddings.py

@@ -1,11 +1,261 @@
+#!/usr/bin/env python3
+"""
+使用BERT + FAISS构建产品描述和Slogan的嵌入数据库
+支持相似性搜索和检索
+"""
+
+import faiss
+import numpy as np
+import pandas as pd
 from sentence_transformers import SentenceTransformer
-import faiss, json, glob, os, numpy as np
-
-model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
-texts=[]; vecs=[]
-for f in glob.glob("nyc_ads_dataset/*.json"):
-    cap=json.load(open(f))["caption"]
-    texts.append(cap); vecs.append(model.encode(cap,normalize_embeddings=True))
-vecs=np.vstack(vecs).astype("float32")
-index=faiss.IndexFlatIP(vecs.shape[1]); index.add(vecs)
-faiss.write_index(index,"prompt.index"); json.dump(texts,open("prompt.txt","w"))
+from datasets import Dataset
+import pickle
+import json
+from typing import List, Dict, Tuple
+import os
+
+class SloganEmbeddingDB:
+    def __init__(self, model_name: str = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"):
+        """
+        初始化BERT+FAISS数据库
+        
+        Args:
+            model_name: 多语言BERT模型，支持中英文
+        """
+        print(f"📥 Loading BERT model: {model_name}")
+        self.model = SentenceTransformer(model_name)
+        self.dimension = self.model.get_sentence_embedding_dimension()
+        
+        # 初始化FAISS索引
+        self.index = faiss.IndexFlatIP(self.dimension)  # 内积相似度
+        self.data = []  # 存储原始数据
+        
+        print(f"✅ Model loaded. Embedding dimension: {self.dimension}")
+    
+    def create_sample_dataset(self) -> Dataset:
+        """创建示例数据集"""
+        sample_data = [
+            # 中文品牌
+            {"business": "肯德基", "category": "快餐", "description": "美式炸鸡快餐连锁", "slogan": "有了肯德基生活好滋味"},
+            {"business": "麦当劳", "category": "快餐", "description": "全球知名汉堡快餐", "slogan": "我就喜欢"},
+            {"business": "星巴克", "category": "咖啡", "description": "全球连锁咖啡店", "slogan": "启发并滋润人类精神"},
+            {"business": "小米", "category": "电子产品", "description": "智能手机和科技产品", "slogan": "让每个人都能享受科技的乐趣"},
+            {"business": "华为", "category": "电子产品", "description": "通信设备和智能手机", "slogan": "构建万物互联的智能世界"},
+            
+            # 英文品牌
+            {"business": "Nike", "category": "运动用品", "description": "Athletic footwear and apparel", "slogan": "Just Do It"},
+            {"business": "Apple", "category": "科技", "description": "Consumer electronics and software", "slogan": "Think Different"},
+            {"business": "Coca-Cola", "category": "饮料", "description": "Carbonated soft drinks", "slogan": "Open Happiness"},
+            {"business": "BMW", "category": "汽车", "description": "Luxury automobiles", "slogan": "The Ultimate Driving Machine"},
+            {"business": "Amazon", "category": "电商", "description": "E-commerce and cloud services", "slogan": "Earth's Most Customer-Centric Company"},
+            
+            # 产品描述
+            {"business": "智能手表", "category": "可穿戴设备", "description": "健康监测和通知功能的智能手表", "slogan": "时刻关注您的健康"},
+            {"business": "电动汽车", "category": "新能源汽车", "description": "零排放环保电动车", "slogan": "绿色出行，智享未来"},
+            {"business": "在线教育平台", "category": "教育科技", "description": "AI驱动的个性化学习平台", "slogan": "让学习更智能"},
+            {"business": "健身APP", "category": "健康应用", "description": "AI私教健身指导应用", "slogan": "随时随地，专业健身"},
+            {"business": "外卖平台", "category": "生活服务", "description": "快速便捷的餐食配送服务", "slogan": "美食到家，生活更美好"},
+        ]
+        
+        return Dataset.from_pandas(pd.DataFrame(sample_data))
+    
+    def build_embeddings(self, dataset: Dataset):
+        """构建嵌入向量并建立FAISS索引"""
+        print("🔨 Building embeddings and FAISS index...")
+        
+        # 准备数据
+        texts = []
+        for item in dataset:
+            # 组合文本：业务名称 + 类别 + 描述
+            combined_text = f"{item['business']} {item['category']} {item['description']}"
+            texts.append(combined_text)
+            
+            # 保存原始数据
+            self.data.append({
+                "business": item["business"],
+                "category": item["category"], 
+                "description": item["description"],
+                "slogan": item["slogan"],
+                "combined_text": combined_text
+            })
+        
+        # 生成嵌入向量
+        print(f"📊 Generating embeddings for {len(texts)} items...")
+        embeddings = self.model.encode(texts, show_progress_bar=True)
+        
+        # 标准化向量（用于余弦相似度）
+        embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
+        
+        # 添加到FAISS索引
+        self.index.add(embeddings.astype('float32'))
+        
+        print(f"✅ Built FAISS index with {self.index.ntotal} vectors")
+    
+    def search_similar(self, query: str, top_k: int = 5) -> List[Dict]:
+        """搜索相似的业务描述"""
+        print(f"🔍 Searching for: '{query}'")
+        
+        # 生成查询向量
+        query_embedding = self.model.encode([query])
+        query_embedding = query_embedding / np.linalg.norm(query_embedding, axis=1, keepdims=True)
+        
+        # FAISS搜索
+        scores, indices = self.index.search(query_embedding.astype('float32'), top_k)
+        
+        # 整理结果
+        results = []
+        for i, (score, idx) in enumerate(zip(scores[0], indices[0])):
+            if idx < len(self.data):
+                result = self.data[idx].copy()
+                result["similarity_score"] = float(score)
+                result["rank"] = i + 1
+                results.append(result)
+        
+        return results
+    
+    def save_database(self, save_path: str = "./slogan_db"):
+        """保存数据库"""
+        os.makedirs(save_path, exist_ok=True)
+        
+        # 保存FAISS索引
+        faiss.write_index(self.index, f"{save_path}/faiss.index")
+        
+        # 保存数据
+        with open(f"{save_path}/data.pkl", "wb") as f:
+            pickle.dump(self.data, f)
+        
+        # 保存配置
+        config = {
+            "model_name": self.model._modules['0'].auto_model.config.name_or_path,
+            "dimension": self.dimension,
+            "total_items": len(self.data)
+        }
+        with open(f"{save_path}/config.json", "w", encoding="utf-8") as f:
+            json.dump(config, f, ensure_ascii=False, indent=2)
+        
+        print(f"💾 Database saved to {save_path}")
+    
+    def load_database(self, load_path: str = "./slogan_db"):
+        """加载数据库"""
+        print(f"📂 Loading database from {load_path}")
+        
+        # 加载FAISS索引
+        self.index = faiss.read_index(f"{load_path}/faiss.index")
+        
+        # 加载数据
+        with open(f"{load_path}/data.pkl", "rb") as f:
+            self.data = pickle.load(f)
+        
+        print(f"✅ Loaded database with {len(self.data)} items")
+    
+    def add_new_item(self, business: str, category: str, description: str, slogan: str):
+        """动态添加新项目"""
+        combined_text = f"{business} {category} {description}"
+        
+        # 生成嵌入
+        embedding = self.model.encode([combined_text])
+        embedding = embedding / np.linalg.norm(embedding, axis=1, keepdims=True)
+        
+        # 添加到索引
+        self.index.add(embedding.astype('float32'))
+        
+        # 添加到数据
+        self.data.append({
+            "business": business,
+            "category": category,
+            "description": description,
+            "slogan": slogan,
+            "combined_text": combined_text
+        })
+        
+        print(f"➕ Added new item: {business}")
+    
+    def generate_slogan_suggestions(self, business_description: str, top_k: int = 3) -> List[str]:
+        """根据业务描述生成Slogan建议"""
+        similar_items = self.search_similar(business_description, top_k)
+        
+        suggestions = []
+        for item in similar_items:
+            suggestions.append({
+                "slogan": item["slogan"],
+                "reference": f"{item['business']} ({item['category']})",
+                "similarity": item["similarity_score"]
+            })
+        
+        return suggestions
+
+def main():
+    """主函数演示"""
+    # 初始化数据库
+    db = SloganEmbeddingDB()
+    
+    # 创建或加载数据
+    if os.path.exists("./slogan_db"):
+        print("📂 Found existing database, loading...")
+        db.load_database()
+    else:
+        print("🆕 Creating new database...")
+        dataset = db.create_sample_dataset()
+        db.build_embeddings(dataset)
+        db.save_database()
+    
+    # 测试搜索
+    test_queries = [
+        "智能穿戴设备健康监测",
+        "环保新能源汽车",
+        "人工智能学习平台",
+        "美式快餐炸鸡",
+        "luxury sports car",
+        "mobile phone technology"
+    ]
+    
+    print("\n" + "="*60)
+    print("🔍 SEARCH RESULTS")
+    print("="*60)
+    
+    for query in test_queries:
+        print(f"\n🔍 Query: {query}")
+        results = db.search_similar(query, top_k=3)
+        
+        for result in results:
+            print(f"  {result['rank']}. {result['business']} ({result['category']})")
+            print(f"     描述: {result['description']}")
+            print(f"     Slogan: {result['slogan']}")
+            print(f"     相似度: {result['similarity_score']:.3f}")
+            print()
+    
+    # 测试Slogan��成建议
+    print("\n" + "="*60)
+    print("💡 SLOGAN SUGGESTIONS")
+    print("="*60)
+    
+    new_business = "AI智能音箱语音助手设备"
+    print(f"\n💡 为 '{new_business}' 生成Slogan建议:")
+    
+    suggestions = db.generate_slogan_suggestions(new_business)
+    for i, suggestion in enumerate(suggestions, 1):
+        print(f"  {i}. \"{suggestion['slogan']}\"")
+        print(f"     参考: {suggestion['reference']}")
+        print(f"     相似度: {suggestion['similarity']:.3f}")
+        print()
+    
+    # 演示动态添加
+    print("\n" + "="*60)
+    print("➕ ADDING NEW ITEM")
+    print("="*60)
+    
+    db.add_new_item(
+        business="智能眼镜",
+        category="AR设备", 
+        description="增强现实智能眼镜产品",
+        slogan="看见未来，触手可及"
+    )
+    
+    # 重新搜索测试
+    print(f"\n🔍 搜索 'AR增强现实产品':")
+    results = db.search_similar("AR增强现实产品", top_k=2)
+    for result in results:
+        print(f"  - {result['business']}: {result['slogan']} (相似度: {result['similarity_score']:.3f})")
+
+if __name__ == "__main__":
+    main()