File size: 7,909 Bytes
5dde370 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 |
import json
import os
import re
import pandas as pd
import nltk
from nltk.corpus import stopwords
from nltk.stem import WordNetLemmatizer
from sklearn.feature_extraction.text import TfidfVectorizer
from typing import List, Dict, Any
from tqdm import tqdm
import hashlib
# 下载必要的NLTK资源
required_packages = ['punkt', 'stopwords', 'wordnet']
for package in required_packages:
try:
nltk.data.find(f'tokenizers/{package}' if package == 'punkt' else f'corpora/{package}')
print(f"Package {package} is already downloaded.")
except LookupError:
print(f"Downloading NLTK {package}...")
nltk.download(package)
class QAProcessor:
def __init__(self, data_dir: str, output_path: str):
self.data_dir = data_dir
self.output_path = output_path
self.lemmatizer = WordNetLemmatizer()
self.stop_words = set(stopwords.words('english'))
self.medical_stopwords = {'disease', 'patient', 'treatment', 'condition', 'symptom', 'doctor', 'health', 'may', 'also', 'one', 'use'}
self.stop_words.update(self.medical_stopwords)
def clean_text(self, text: str) -> str:
"""清理文本"""
if isinstance(text, str):
# 移除HTML标签
text = re.sub(r'<.*?>', ' ', text)
# 移除特殊字符但保留字母和数字
text = re.sub(r'[^\w\s]', ' ', text)
# 移除多余的空格
text = re.sub(r'\s+', ' ', text).strip().lower()
return text
return ""
def simple_tokenize(self, text: str) -> List[str]:
"""简单的分词函数"""
text = text.lower().strip()
tokens = re.findall(r'\b\w+\b', text)
return tokens
def extract_keywords(self, text: str, top_n: int = 10) -> List[str]:
"""使用TF-IDF提取关键词"""
if not isinstance(text, str) or not text.strip():
return []
tokens = self.simple_tokenize(text)
filtered_tokens = [self.lemmatizer.lemmatize(token)
for token in tokens
if token.isalpha() and token not in self.stop_words and len(token) > 2]
if filtered_tokens:
vectorizer = TfidfVectorizer(max_features=top_n)
try:
tfidf_matrix = vectorizer.fit_transform([' '.join(filtered_tokens)])
feature_names = vectorizer.get_feature_names_out()
scores = zip(feature_names, tfidf_matrix.toarray()[0])
sorted_scores = sorted(scores, key=lambda x: x[1], reverse=True)
return [word for word, score in sorted_scores[:top_n] if score > 0]
except Exception as e:
print(f"TF-IDF提取失败: {e}")
return filtered_tokens[:top_n]
return []
def process_data(self) -> tuple[List[Dict[str, Any]], Dict[str, List[str]]]:
"""处理所有数据源"""
qa_database = []
keyword_index = {} # 倒排索引：keyword → [qa_id1, qa_id2, ...]
# 处理Healthline文章
healthline_path = os.path.join(self.data_dir,'Healthline', 'healthline_articles_text.csv')
if os.path.exists(healthline_path):
print("处理Healthline数据...")
healthline_df = pd.read_csv(healthline_path)
for idx, row in tqdm(healthline_df.iterrows(), total=len(healthline_df)):
title = row.get('title', '')
content = row.get('content', '')
if not isinstance(title, str) or not isinstance(content, str):
continue
clean_title = self.clean_text(title)
clean_content = self.clean_text(content)
qa_pair = {
'id': f"healthline_{idx}",
'source': 'healthline',
'question': clean_title,
'answer': clean_content,
'keywords': self.extract_keywords(clean_title + " " + clean_content)
}
qa_database.append(qa_pair)
# 更新倒排索引
for keyword in qa_pair['keywords']:
if keyword not in keyword_index:
keyword_index[keyword] = []
keyword_index[keyword].append(qa_pair['id'])
# 处理MedQA数据
medqa_dir = os.path.join(self.data_dir, 'MedQA')
if os.path.exists(medqa_dir):
print("处理MedQA数据...")
for file_name in os.listdir(medqa_dir):
if file_name.endswith('.csv'):
dataset_name = file_name.split('.')[0]
df = pd.read_csv(os.path.join(medqa_dir, file_name))
question_col = next((col for col in ['Question', 'question'] if col in df.columns), None)
answer_col = next((col for col in ['Answer', 'answer'] if col in df.columns), None)
if not question_col or not answer_col:
print(f"跳过 {dataset_name} - 缺少问题/答案列")
continue
for idx, row in tqdm(df.iterrows(), total=len(df)):
question = row.get(question_col, '')
answer = row.get(answer_col, '')
if not isinstance(question, str) or not isinstance(answer, str):
continue
clean_question = self.clean_text(question)
clean_answer = self.clean_text(answer)
qa_pair = {
'id': f"{dataset_name}_{idx}",
'source': dataset_name,
'question': clean_question,
'answer': clean_answer,
'keywords': self.extract_keywords(clean_question + " " + clean_answer)
}
qa_database.append(qa_pair)
for keyword in qa_pair['keywords']:
if keyword not in keyword_index:
keyword_index[keyword] = []
keyword_index[keyword].append(qa_pair['id'])
return qa_database, keyword_index
def save_results(self, qa_database: List[Dict[str, Any]], keyword_index: Dict[str, List[str]]):
"""保存处理后的数据"""
qa_output = os.path.join(self.output_path, 'cleaned_qa','qa_database.json')
keyword_output = os.path.join(self.output_path, 'keywords','keyword_index.json')
os.makedirs(os.path.dirname(qa_output), exist_ok=True)
os.makedirs(os.path.dirname(keyword_output), exist_ok=True)
with open(qa_output, 'w', encoding='utf-8') as f:
json.dump(qa_database, f, ensure_ascii=False, indent=2)
with open(keyword_output, 'w', encoding='utf-8') as f:
json.dump(keyword_index, f, ensure_ascii=False, indent=2)
print(f"处理完成的QA对数量: {len(qa_database)}")
print(f"关键词索引中的关键词数量: {len(keyword_index)}")
print(f"数据已保存到: {self.output_path}")
if __name__ == "__main__":
# 设置路径
data_dir = './Data/raw'
output_path = './Data/Processed/'
# 创建处理器实例并处理数据
processor = QAProcessor(data_dir, output_path)
qa_database, keyword_index = processor.process_data()
processor.save_results(qa_database, keyword_index) |