Update app.py

app.py

@@ -1,4 +1,4 @@
-from datasets import load_dataset
+import pandas as pd
 import numpy as np
 from transformers import AutoTokenizer, AutoModel, pipeline
 import torch
@@ -7,67 +7,96 @@ from sentence_transformers import SentenceTransformer
 from sklearn.metrics.pairwise import cosine_similarity
 import gradio as gr
 
-# 加载模型（全局加载提升性能）
-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)
+# 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
 
-# 加载bookcorpus数据集
+# Load book data
 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]  # 过滤短文本
+    # 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()
+    return books
 
-# 文本嵌入生成（复用原始代码逻辑）
+# Mean pooling for sentence embeddings
 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)
 
-def get_embeddings(texts):
-    encoded_input = TOKENIZER(texts, padding=True, truncation=True, return_tensors='pt')
+# 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)
+        model_output = model(**encoded_input)
     embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
-    return F.normalize(embeddings, p=2, dim=1)
+    embeddings = F.normalize(embeddings, p=2, dim=1)
+    return embeddings
 
-# 相似度计算适配新数据集
-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])
+# 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)
+    
+    # Get embeddings for book titles and descriptions
+    book_texts = books['title'] + " " + books['author']
+    book_embeddings = get_embeddings(book_texts.tolist(), tokenizer, model)
+    
+    # Calculate similarity
     similarities = cosine_similarity(keyword_embedding, book_embeddings)[0]
+    
+    # Get top matches
     top_indices = np.argsort(similarities)[-top_k:][::-1]
-    return [books[i] for i in top_indices]
+    results = books.iloc[top_indices].copy()
+    results['similarity'] = similarities[top_indices]
+    
+    return results
 
-# 摘要生成适配长文本
-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
+# 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()
-    similar_books = find_similar_books(keywords, books)
     
+    # Find similar books
+    similar_books = find_similar_books(keywords, books, tokenizer, model)
+    
+    # Generate output
     output = []
-    for i, book in enumerate(similar_books, 1):
-        summary = summarize_description(book["text"])
-        output.append(f"{i}. {summary}\n")
+    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")
+    
     return "\n".join(output)
 
-# Gradio界面保持相同
+# Gradio interface
 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"
+    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."
 )
 
 if __name__ == "__main__":