|
|
|
|
| """
|
| Emotion Summary Model - 推理示例
|
| """
|
|
|
| from transformers import MT5ForConditionalGeneration, MT5Tokenizer
|
| import json
|
| import torch
|
|
|
| def load_model(model_path="./emotion_summary"):
|
| """加载模型和tokenizer"""
|
| print(f"Loading model from {model_path}...")
|
| model = MT5ForConditionalGeneration.from_pretrained(model_path)
|
| tokenizer = MT5Tokenizer.from_pretrained(model_path)
|
|
|
|
|
| device = "cuda" if torch.cuda.is_available() else "cpu"
|
| model = model.to(device)
|
| model.eval()
|
|
|
| print(f"Model loaded on {device}")
|
| return model, tokenizer, device
|
|
|
| def summarize_case(model, tokenizer, device, case_data, field="cause"):
|
| """
|
| 对案例进行总结
|
|
|
| Args:
|
| model: 模型
|
| tokenizer: tokenizer
|
| device: 设备
|
| case_data: 案例数据（字典）
|
| field: 要生成的字段 (cause/symptoms/treatment_process/illness_characteristics/treatment_effect)
|
|
|
| Returns:
|
| 生成的总结文本
|
| """
|
|
|
|
|
| case_desc = " ".join(case_data.get("case_description", []))
|
| consultation = " ".join(case_data.get("consultation_process", []))
|
| reflection = case_data.get("experience_and_reflection", "")
|
|
|
| full_text = f"Case: {case_desc}\nConsultation: {consultation}\nReflection: {reflection}"
|
|
|
|
|
| prompts = {
|
| "cause": f"Extract cause from: {full_text}",
|
| "symptoms": f"Extract symptoms from: {full_text}",
|
| "treatment_process": f"Extract treatment process from: {full_text}",
|
| "illness_characteristics": f"Extract illness characteristics from: {full_text}",
|
| "treatment_effect": f"Extract treatment effect from: {full_text}"
|
| }
|
|
|
| input_text = prompts.get(field, full_text)
|
|
|
|
|
| input_ids = tokenizer.encode(
|
| input_text,
|
| return_tensors="pt",
|
| max_length=512,
|
| truncation=True
|
| ).to(device)
|
|
|
|
|
| with torch.no_grad():
|
| output_ids = model.generate(
|
| input_ids,
|
| max_length=256,
|
| num_beams=4,
|
| early_stopping=True,
|
| no_repeat_ngram_size=3
|
| )
|
|
|
|
|
| output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
|
| return output_text
|
|
|
| def process_test_file(input_file, output_file, model_path="./emotion_summary"):
|
| """处理测试文件"""
|
|
|
|
|
| model, tokenizer, device = load_model(model_path)
|
|
|
|
|
| test_data = []
|
| with open(input_file, 'r', encoding='utf-8') as f:
|
| for line in f:
|
| if line.strip():
|
| test_data.append(json.loads(line))
|
|
|
| print(f"\nProcessing {len(test_data)} samples...")
|
|
|
| results = []
|
| for i, sample in enumerate(test_data, 1):
|
| print(f" [{i}/{len(test_data)}] Processing ID: {sample['id']}...")
|
|
|
| result = {
|
| "id": sample["id"],
|
| "predicted_cause": summarize_case(model, tokenizer, device, sample, "cause"),
|
| "predicted_symptoms": summarize_case(model, tokenizer, device, sample, "symptoms"),
|
| "predicted_treatment_process": summarize_case(model, tokenizer, device, sample, "treatment_process"),
|
| "predicted_illness_Characteristics": summarize_case(model, tokenizer, device, sample, "illness_characteristics"),
|
| "predicted_treatment_effect": summarize_case(model, tokenizer, device, sample, "treatment_effect")
|
| }
|
|
|
| results.append(result)
|
|
|
|
|
| with open(output_file, 'w', encoding='utf-8') as f:
|
| for result in results:
|
| json.dump(result, f, ensure_ascii=False)
|
| f.write('\n')
|
|
|
| print(f"\n✓ Results saved to {output_file}")
|
|
|
| if __name__ == "__main__":
|
|
|
| sample_case = {
|
| "id": 1,
|
| "case_description": ["A 34-year-old male with health anxiety..."],
|
| "consultation_process": ["The consultation began with..."],
|
| "experience_and_reflection": "This case demonstrates..."
|
| }
|
|
|
| model, tokenizer, device = load_model()
|
|
|
| print("\nGenerating summaries...")
|
| cause = summarize_case(model, tokenizer, device, sample_case, "cause")
|
| print(f"\nCause: {cause}")
|
|
|
|
|
|
|
|
|
