Create app

app

@@ -0,0 +1,91 @@
+# 导入必要库
+from datasets import load_dataset
+import pandas as pd
+import torch
+from sentence_transformers import SentenceTransformer, util
+from transformers import pipeline
+
+# ----------------------
+# 1. 加载数据集
+# ----------------------
+def load_book_data():
+    # 加载 bookcorpus 数据集（仅保留标题和摘要）
+    dataset = load_dataset("bookcorpus", split="train")
+    books = pd.DataFrame(dataset)[["title", "text"]].rename(columns={"text": "description"})
+    
+    # 过滤空值并截断长文本（可选）
+    books = books.dropna().head(1000)  # 取前1000条数据便于演示
+    books["description"] = books["description"].apply(lambda x: x[:5000])  # 截断至5000字以内
+    return books
+
+# ----------------------
+# 2. 初始化模型
+# ----------------------
+def initialize_models():
+    # 语义搜索模型
+    embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+    
+    # 摘要生成模型
+    summarizer = pipeline(
+        "summarization",
+        model="facebook/bart-large-cnn",
+        max_length=150,
+        min_length=30,
+        do_sample=False
+    )
+    return embedder, summarizer
+
+# ----------------------
+# 3. 关键词搜索与推荐
+# ----------------------
+def search_similar_books(keywords, books, embedder, top_k=5):
+    # 生成关键词嵌入
+    keyword_embedding = embedder.encode(keywords, convert_to_tensor=True)
+    
+    # 生成书籍嵌入（批量处理）
+    book_embeddings = torch.stack([
+        embedder.encode(title + " " + desc, convert_to_tensor=True) 
+        for title, desc in zip(books["title"], books["description"])
+    ])
+    
+    # 计算余弦相似度
+    cos_scores = util.cos_sim(keyword_embedding, book_embeddings)[0]
+    
+    # 获取 top-k 结果
+    top_results = torch.topk(cos_scores, k=top_k).indices.tolist()
+    return books.iloc[top_results]
+
+# ----------------------
+# 4. 生成摘要并输出
+# ----------------------
+def generate_book_summaries(books, summarizer):
+    results = []
+    for idx, row in books.iterrows():
+        summary = summarizer(row["description"], max_length=150)[0]["summary_text"]
+        results.append({
+            "title": row["title"],
+            "summary": summary,
+            "similarity": "{:.2f}".format(float(cos_scores[idx]))  # 可选：添加相似度分数
+        })
+    return results
+
+# ----------------------
+# 5. 主函数与交互
+# ----------------------
+if __name__ == "__main__":
+    # 加载数据与模型
+    books = load_book_data()
+    embedder, summarizer = initialize_models()
+    
+    # 用户输入关键词
+    user_keywords = "fantasy adventure magic"  # 示例关键词，可替换为用户输入
+    
+    # 执行搜索与摘要生成
+    similar_books = search_similar_books(user_keywords, books, embedder)
+    summaries = generate_book_summaries(similar_books, summarizer)
+    
+    # 打印结果
+    for i, book in enumerate(summaries, 1):
+        print(f"📚 Book {i}: {book['title']}")
+        print(f"🌟 Similarity: {book['similarity']}")
+        print(f"📝 Summary: {book['summary']}\n")