Update app.py

app.py

@@ -1,103 +1,126 @@
-import pandas as pd
+# app.py
+from datasets import load_dataset
 import numpy as np
-from transformers import AutoTokenizer, AutoModel, pipeline
 import torch
 import torch.nn.functional as F
-from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModel, pipeline
 from sklearn.metrics.pairwise import cosine_similarity
 import gradio as gr
+import re
 
-# Load models
-def load_models():
-    # For semantic search
-    model_name = "sentence-transformers/all-MiniLM-L6-v2"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModel.from_pretrained(model_name)
-    
-    # For summarization
-    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
-    
-    return tokenizer, model, summarizer
+# 全局配置
+MODEL_NAME = "sentence-transformers/all-mpnet-base-v2"  # 更强大的语义模型
+SUMMARIZER_NAME = "facebook/bart-large-cnn"
+DATASET_NAME = "bookcorpus"
+CACHE_DIR = "./data-cache"
+
+# 预加载资源
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModel.from_pretrained(MODEL_NAME)
+summarizer = pipeline("summarization", SUMMARIZER_NAME)
 
-# Load book data
-def load_data():
-    # Load the Goodreads dataset (adjust path as needed)
-    books = pd.read_csv("bookcorpus.csv")
-    # Keep only relevant columns and drop rows with missing descriptions
-    books = books[['title', 'author']].dropna()
+# 加载并预处理书籍数据
+def load_books():
+    dataset = load_dataset(DATASET_NAME, split='train', streaming=True)
+    books = []
+    for book in dataset.take(50000):  # 取5万本书
+        text = book['text'].strip()
+        if len(text) > 500:  # 过滤短文本
+            title = re.findall(r'"([^"]*)"', text[:200])  # 尝试提取标题
+            books.append({
+                "text": text,
+                "title": title[0] if title else "Untitled Book"
+            })
     return books
 
-# Mean pooling for sentence embeddings
+# 生成语义嵌入
+def get_embeddings(texts):
+    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt", max_length=512)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    embeddings = mean_pooling(outputs, inputs['attention_mask'])
+    return F.normalize(embeddings, p=2, dim=1)
+
+# 平均池化
 def mean_pooling(model_output, attention_mask):
-    token_embeddings = model_output[0]
+    token_embeddings = model_output.last_hidden_state
     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)
+    return torch.sum(token_embedding * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
 
-# Get embeddings for text
-def get_embeddings(texts, tokenizer, model):
-    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'])
-    embeddings = F.normalize(embeddings, p=2, dim=1)
-    return embeddings
+# 智能摘要生成
+def generate_summary(text):
+    inputs = tokenizer(
+        "summarize: " + text,
+        max_length=1024,
+        truncation=True,
+        return_tensors="pt"
+    )
+    summary_ids = summarizer.model.generate(
+        inputs.input_ids,
+        max_length=150,
+        min_length=50,
+        length_penalty=2.0,
+        num_beams=4,
+        early_stopping=True
+    )
+    return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
 
-# Find most similar books
-def find_similar_books(keywords, books, tokenizer, model, top_k=5):
-    # Get embeddings for keywords
-    keyword_embedding = get_embeddings(keywords, tokenizer, model).mean(0).unsqueeze(0)
+# 核心推荐逻辑
+def recommend_books(keywords, top_k=5):
+    # 清洗输入
+    keywords = re.sub(r'[^\w\s,]', '', keywords).lower()
+    keywords = [k.strip() for k in keywords.split(',') if k.strip()]
     
-    # Get embeddings for book titles and descriptions
-    book_texts = books['title'] + " " + books['author']
-    book_embeddings = get_embeddings(book_texts.tolist(), tokenizer, model)
+    if len(keywords) < 2:
+        return "❗ Please enter at least 2 keywords (e.g. 'fantasy, magic')"
     
-    # Calculate similarity
-    similarities = cosine_similarity(keyword_embedding, book_embeddings)[0]
+    # 获取嵌入
+    keyword_emb = get_embeddings([" ".join(keywords)]).mean(dim=0)
+    book_embs = get_embeddings([f"{b['title']} {b['text']}" for b in books])
     
-    # Get top matches
-    top_indices = np.argsort(similarities)[-top_k:][::-1]
-    results = books.iloc[top_indices].copy()
-    results['similarity'] = similarities[top_indices]
+    # 计算相似度
+    sim_scores = cosine_similarity(keyword_emb.reshape(1,-1), book_embs)[0]
+    top_indices = np.argsort(sim_scores)[-top_k:][::-1]
     
+    # 生成结果
+    results = []
+    for idx in top_indices:
+        book = books[idx]
+        summary = generate_summary(book['text'])
+        results.append({
+            "title": book['title'],
+            "summary": summary,
+            "score": f"{sim_scores[idx]:.2f}"
+        })
     return results
 
-# Summarize book description
-def summarize_description(description, summarizer):
-    if len(description.split()) > 100:  # Only summarize long descriptions
-        summary = summarizer(description, max_length=130, min_length=30, do_sample=False)
-        return summary[0]['summary_text']
-    return description
-
-# Main function
-def recommend_books(keywords):
-    # Split keywords by comma or space
-    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."
-    
-    # Load models and data
-    tokenizer, model, summarizer = load_models()
-    books = load_data()
+# Gradio界面
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 📚 智能图书推荐系统")
     
-    # Find similar books
-    similar_books = find_similar_books(keywords, books, tokenizer, model)
+    with gr.Row():
+        inputs = gr.Textbox(label="输入关键词（用逗号分隔）", placeholder="例如：sci-fi, time travel")
+        outputs = gr.JSON(label="推荐结果")
     
-    # Generate output
-    output = []
-    for i, (_, row) in enumerate(similar_books.iterrows(), 1):
-        summary = summarize_description(row['description'], summarizer)
-        output.append(f"{i}. {row['title']}\n   Summary: {summary}\n")
+    examples = gr.Examples(
+        examples=[
+            ["romance, paris"],
+            ["mystery, detective"],
+            ["science fiction, space opera"]
+        ],
+        inputs=[inputs]
+    )
     
-    return "\n".join(output)
+    inputs.submit(
+        fn=recommend_books,
+        inputs=inputs,
+        outputs=outputs
+    )
 
-# Gradio interface
-iface = gr.Interface(
-    fn=recommend_books,
-    inputs=gr.Textbox(label="Enter at least 3 keywords (comma or space separated)"),
-    outputs=gr.Textbox(label="Recommended Books"),
-    title="Book Recommendation Engine",
-    description="Enter 3 or more keywords to find relevant books and get summaries of their plots."
-)
+# 初始化数据
+print("Loading book data...")
+books = load_books()
+print(f"Loaded {len(books)} books")
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860)