File size: 4,679 Bytes
d22f4ca | 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 | #!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
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)
# 如果有GPU就使用GPU
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}"
# 根据字段构建不同的prompt
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}")
# 如果要处理整个测试文件，取消下面的注释：
# process_test_file("data/test/Emotion_Summary.jsonl", "results/predictions.jsonl")
|