Back-end/data_processing/process_qa.py

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)