Update app.py

app.py

@@ -1,91 +1,65 @@
-# 导入必要库
-from datasets import load_dataset
-import pandas as pd
-import torch
-from sentence_transformers import SentenceTransformer, util
-from transformers import pipeline
+# 加载模型（全局加载提升性能）
+MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+SUMMARIZER = pipeline("summarization", model="facebook/bart-large-cnn")
+TOKENIZER = AutoTokenizer.from_pretrained(MODEL_NAME)
+MODEL = AutoModel.from_pretrained(MODEL_NAME)
 
-# ----------------------
-# 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
+# 加载bookcorpus数据集
+def load_data():
+    dataset = load_dataset("bookcorpus", streaming=True)  # 启用流式读取
+    books = dataset["train"].take(100_000)  # 取前10万条数据
+    return [{"text": x["text"]} for x in books if len(x["text"]) > 100]  # 过滤短文本
 
-# ----------------------
-# 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
+# 文本嵌入生成（复用原始代码逻辑）
+def mean_pooling(model_output, attention_mask):
+    token_embeddings = model_output[0]
+    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
 
-# ----------------------
-# 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"])
-    ])
+def get_embeddings(texts):
+    encoded_input = TOKENIZER(texts, padding=True, truncation=True, return_tensors='pt')
+    with torch.no_grad():
+        model_output = MODEL(**encoded_input)
+    embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
+    return F.normalize(embeddings, p=2, dim=1)
+
+# 相似度计算适配新数据集
+def find_similar_books(keywords, books, top_k=5):
+    keyword_embedding = get_embeddings(keywords).mean(0).unsqueeze(0)
+    book_embeddings = get_embeddings([book["text"] for book in books])
+    similarities = cosine_similarity(keyword_embedding, book_embeddings)[0]
+    top_indices = np.argsort(similarities)[-top_k:][::-1]
+    return [books[i] for i in top_indices]
+
+# 摘要生成适配长文本
+def summarize_description(text):
+    if len(text.split()) > 500:
+        return SUMMARIZER(text, max_length=150, min_length=50, do_sample=False)[0]['summary_text']
+    return text
+
+# 主逻辑流程
+def recommend_books(keywords):
+    keywords = [kw.strip() for kw in keywords.replace(',', ' ').split() if kw.strip()]
+    if len(keywords) < 3:
+        return "Please enter at least 3 keywords separated by commas or spaces."
     
-    # 计算余弦相似度
-    cos_scores = util.cos_sim(keyword_embedding, book_embeddings)[0]
+    books = load_data()
+    similar_books = find_similar_books(keywords, books)
     
-    # 获取 top-k 结果
-    top_results = torch.topk(cos_scores, k=top_k).indices.tolist()
-    return books.iloc[top_results]
+    output = []
+    for i, book in enumerate(similar_books, 1):
+        summary = summarize_description(book["text"])
+        output.append(f"{i}. {summary}\n")
+    return "\n".join(output)
 
-# ----------------------
-# 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
+# Gradio界面保持相同
+iface = gr.Interface(
+    fn=recommend_books,
+    inputs=gr.Textbox(label="Enter 3+ keywords (comma/space separated)"),
+    outputs=gr.Textbox(label="Recommended Book Passages"),
+    title="Book Corpus Semantic Search",
+    description="Search through 100,000 book passages from bookcorpus dataset"
+)
 
-# ----------------------
-# 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")
+    iface.launch()