Update app.py

app.py

@@ -1,103 +1,67 @@
-import pandas as pd
+# 安装依赖
+!pip install datasets sentence-transformers transformers torch
+
+# 导入库
+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 sklearn.metrics.pairwise import cosine_similarity
-import gradio as gr
+from sentence_transformers import SentenceTransformer, util
+from transformers import pipeline
 
-# Load models
-def load_models():
-    # For semantic search
-    model_name = "sentence-transformers/all-mpnet-base-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
+# 加载数据集
+dataset = load_dataset("Pradeep016/career-guidance-qa-dataset", split="train")
+# 过滤无效数据（确保question和answer非空）
+dataset = dataset.filter(lambda x: x["question"] and x["answer"])
 
-# Load book data
-def load_data():
-    # Load the Goodreads dataset (adjust path as needed)
-    books = pd.read_csv("books.csv")
-    # Keep only relevant columns and drop rows with missing descriptions
-    books = books[['title', 'description']].dropna()
-    return books
+# 构建职位知识库（职位名称 + 问题-答案对）
+def build_knowledge_base(dataset):
+    knowledge_base = []
+    for item in dataset:
+        role = item["role"]
+        question = item["question"]
+        answer = item["answer"]
+        # 合并职位名称与问题，增强语义关联
+        entry = f"{role} | {question}: {answer}"
+        knowledge_base.append(entry)
+    return knowledge_base
 
-# 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)
+knowledge_base = build_knowledge_base(dataset)
 
-# 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
+# 初始化语义搜索模型
+embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+# 预计算知识库嵌入向量
+knowledge_embeddings = embedder.encode(knowledge_base, convert_to_tensor=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)
-    
-    # Get embeddings for book titles and descriptions
-    book_texts = books['title'] + " " + books['description']
-    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]
-    results = books.iloc[top_indices].copy()
-    results['similarity'] = similarities[top_indices]
+def career_qa(user_input):
+    # 1. 语义搜索匹配相关职位
+    input_embedding = embedder.encode(user_input, convert_to_tensor=True)
+    # 计算余弦相似度
+    cos_scores = util.cos_sim(input_embedding, knowledge_embeddings)[0]
+    # 取前3个最相关条目
+    top_indices = np.argsort(cos_scores)[-3:][::-1]
+    top_matches = [knowledge_base[idx] for idx in top_indices]
     
+    # 2. 从匹配条目中提取答案
+    qa_pipeline = pipeline("question-answering", model="distilbert-base-uncased-finetuned-squad2")
+    results = []
+    for match in top_matches:
+        role = match.split(" | ")[0]
+        context = match.split(" | ")[1]
+        # 固定问题为“请介绍这个职位”
+        result = qa_pipeline(question="请介绍这个职位", context=context)
+        results.append({
+            "职位名称": role,
+            "简介": result["answer"],
+            "置信度": result["score"]
+        })
     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()
-    
-    # Find similar books
-    similar_books = find_similar_books(keywords, books, tokenizer, model)
-    
-    # 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")
-    
-    return "\n".join(output)
+# 用户输入职业关键词
+user_query = "零售经理"
+results = career_qa(user_query)
 
-# 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."
-)
+# 输出结果
+for res in results:
+    print(f"职位：{res['职位名称']}")
+    print(f"简介：{res['简介']}")
+    print(f"置信度：{res['置信度']:.2f}\n")
 
-if __name__ == "__main__":
-    iface.launch()