File size: 3,364 Bytes
60b21d3 | 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 | # SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
#
# SPDX-License-Identifier: MIT
import re
import sys
from typing import Dict, Any
from common_evaluator import CommonEvaluator
from opentslm.time_series_datasets.sleep.SleepEDFCoTQADataset import SleepEDFCoTQADataset
def extract_sleep_stage_from_prediction(prediction: str) -> str:
"""
Extract the sleep stage label from the model's prediction.
- If 'Answer:' is present, take everything after the last 'Answer:'
- Otherwise, take the last word/phrase
- Strips whitespace and punctuation
"""
pred = prediction.strip()
# Find the last occurrence of 'Answer:' (case-insensitive)
match = list(re.finditer(r'answer:\s*', pred, re.IGNORECASE))
if match:
# Take everything after the last 'Answer:'
start = match[-1].end()
label = pred[start:].strip()
else:
# Take the last word/phrase (sleep stages can be multi-word)
words = pred.split()
if len(words) >= 2 and words[-2].lower() in ['non-rem', 'rem']:
# Handle multi-word sleep stages like "Non-REM stage 1", "REM sleep"
label = ' '.join(words[-2:])
else:
# Single word stage
label = words[-1] if words else ''
# Look for "Answer: [answer]" pattern with more precise matching
answer_match = re.search(r'Answer:\s*(.+?)(?:\.|$)', pred, re.IGNORECASE)
if answer_match:
label = answer_match.group(1).strip()
# Remove trailing period if present
if label.endswith('.'):
label = label[:-1]
return label.strip().lower()
# Remove trailing punctuation (e.g., period, comma)
label = re.sub(r'[\.,;:!?]+$', '', label)
return label.strip().lower()
def evaluate_sleep_cot(ground_truth: str, prediction: str) -> Dict[str, Any]:
"""
Evaluate SleepEDFCoTQADataset predictions against ground truth.
Extracts the sleep stage label from the end of the model's output and compares to ground truth.
"""
gt_clean = ground_truth.lower().strip()
pred_label = extract_sleep_stage_from_prediction(prediction)
accuracy = int(gt_clean == pred_label)
return {"accuracy": accuracy}
def main():
"""Main function to run SleepEDF CoT evaluation."""
if len(sys.argv) != 2:
print("Usage: python evaluate_sleep_cot.py <model_name>")
print("Example: python evaluate_sleep_cot.py meta-llama/Llama-3.2-1B")
sys.exit(1)
model_name = sys.argv[1]
dataset_classes = [SleepEDFCoTQADataset]
evaluation_functions = {
"SleepEDFCoTQADataset": evaluate_sleep_cot,
}
evaluator = CommonEvaluator()
results_df = evaluator.evaluate_multiple_models(
model_names=[model_name],
dataset_classes=dataset_classes,
evaluation_functions=evaluation_functions,
max_samples=None, # Set to None for full evaluation
max_new_tokens=1000, # Sleep stage classification might need more tokens for reasoning
)
print("\n" + "="*80)
print("FINAL RESULTS SUMMARY")
print("="*80)
print(results_df.to_string(index=False))
return results_df
if __name__ == "__main__":
main() |