File size: 2,745 Bytes
e66f29b |
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 |
#!/usr/bin/env python3
"""
OmniCoreX Evaluation CLI Script
Performs evaluation on a validation dataset, computing relevant metrics.
Supports loading checkpoints and configurable evaluation batch size.
"""
import argparse
import os
import torch
from utils import load_config_file, setup_logging
from model import OmniCoreXModel
from data_loader import create_omncorex_dataloader
def parse_args():
parser = argparse.ArgumentParser(description="Evaluate OmniCoreX Model")
parser.add_argument("--config", type=str, default="config.yaml", help="Path to config file")
parser.add_argument("--checkpoint", type=str, required=True, help="Path to model checkpoint")
parser.add_argument("--batch_size", type=int, help="Batch size for evaluation")
parser.add_argument("--log_file", type=str, help="Log file path")
return parser.parse_args()
def main():
args = parse_args()
config = load_config_file(args.config)
logger = setup_logging(log_file=args.log_file)
logger.info(f"Loading checkpoint from {args.checkpoint}")
data_cfg = config["data"]
eval_batch_size = args.batch_size or data_cfg.get("batch_size", 16)
valid_loader = create_omncorex_dataloader(
metadata_path=data_cfg["metadata_path"],
modalities=data_cfg["modalities"],
tokenizer=None,
batch_size=eval_batch_size,
shuffle=False,
num_workers=data_cfg.get("num_workers", 4),
augmentation=False
)
model_cfg = config["model"]
model = OmniCoreXModel(
stream_configs=model_cfg["streams"],
embed_dim=model_cfg.get("architecture", {}).get("embed_dim", 768),
num_layers=model_cfg.get("architecture", {}).get("num_layers", 24),
num_heads=model_cfg.get("architecture", {}).get("num_heads", 12),
dropout=model_cfg.get("architecture", {}).get("dropout", 0.1)
)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
checkpoint = torch.load(args.checkpoint, map_location=device)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()
total_loss = 0.0
count = 0
criterion = torch.nn.CrossEntropyLoss()
with torch.no_grad():
for batch in valid_loader:
inputs = {k: v.to(device) if torch.is_tensor(v) else v for k, v in batch.items()}
outputs = model(**inputs)
if "labels" in inputs:
loss = criterion(outputs.view(-1, outputs.size(-1)), inputs["labels"].view(-1))
total_loss += loss.item()
count += 1
avg_loss = total_loss / count if count > 0 else 0.0
logger.info(f"Evaluation completed. Average loss: {avg_loss:.6f}")
if __name__ == "__main__":
main()
|