code: def main(
      batch_size: int = 128,
      weak_model_name: str = "alexnet",
      strong_model_name: str = "resnet50_dino",
      n_train: int = 40000,
      seed: int = 0,
      data_path: str = "/root/",
      n_epochs: int = 10,
      lr: float = 1e-3,
  ):
      weak_model = get_model(weak_model_name)
      strong_model = get_model(strong_model_name)
      _, loader = get_imagenet(
          data_path, split="val", batch_size=batch_size, shuffle=False
      )
      print("Getting weak labels...")
      _, gt_labels, weak_labels, weak_acc = get_embeddings(
          weak_model, loader
      )
      print(f"Weak label accuracy: {weak_acc:.3f}")
      print("Getting strong embeddings...")
      embeddings, strong_gt_labels, _, _ = get_embeddings(
          strong_model, loader
      )
      assert torch.all(gt_labels == strong_gt_labels)
      del strong_gt_labels
  
      order = np.arange(len(embeddings))
      rng = np.random.default_rng(seed)
      rng.shuffle(order)
      x = embeddings[order]
      y = gt_labels[order]
      yw = weak_labels[order]
      x_train, x_test = x[:n_train], x[n_train:]
      y_train, y_test = y[:n_train], y[n_train:]
      yw_train, yw_test = yw[:n_train], yw[n_train:]
      yw_test = torch.argmax(yw_test, dim=1)
      eval_datasets = {
          "test": (x_test, y_test),
          "test_weak": (x_test, yw_test),
      }
  
      print("Training logreg on weak labels...")
      results_weak = train_logreg(
          x_train, yw_train, eval_datasets, n_epochs=n_epochs, lr=lr
      )
      print(f"Final accuracy: {results_weak['test']:.3f}")
      print(
          f"Final supervisor-student agreement: {results_weak['test_weak']:.3f}"
      )
      print(
          f"Accuracy by epoch: {[acc.item() for acc in results_weak['test_all']]}"
      )
      print(
          f"Supervisor-student agreement by epoch: {[acc.item() for acc in results_weak['test_weak_all']]}"
      )
  
      print("Training logreg on ground truth labels...")
      results_gt = train_logreg(
          x_train, y_train, eval_datasets, n_epochs=n_epochs, lr=lr
      )
      print(f"Final accuracy: {results_gt['test']:.3f}")
      print(
          f"Accuracy by epoch: {[acc.item() for acc in results_gt['test_all']]}"
      )
  
      print("\n\n" + "=" * 100)
      print(f"Weak label accuracy: {weak_acc:.3f}")
      print(f"Weak→Strong accuracy: {results_weak['test']:.3f}")
      print(f"Strong accuracy: {results_gt['test']:.3f}")
      print("=" * 100)
  
file_path: weak-to-strong/vision/run_weak_strong.py
imports: import fire
  import numpy
  import torch
  import tqdm
  from data import get_imagenet
  from models import alexnet, resnet50_dino, vitb8_dino
  from torch import nn
label_desc: Use image node_id 4 for calling a global function with example usage: main(batch_size, weak_model_name, strong_model_name, n_train, seed, data_path, n_epochs, lr) without return types