Add self-contained training script

train_job.py

@@ -0,0 +1,576 @@
+#!/usr/bin/env python3
+"""
+GeneSetCLIP Training Job — Self-contained script for HF Jobs.
+
+Downloads MSigDB data from Hub, then trains the contrastive model.
+"""
+
+import os
+import sys
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+# Step 1: Download data from Hub
+print("=" * 70)
+print("Step 1: Downloading MSigDB data from Hub...")
+print("=" * 70)
+
+from huggingface_hub import hf_hub_download
+
+DATA_DIR = "/tmp/data"
+os.makedirs(DATA_DIR, exist_ok=True)
+
+for split in ["train", "val", "test"]:
+    path = hf_hub_download(
+        repo_id="AliSaadatV/msigdb-contrastive-data",
+        filename=f"{split}.jsonl",
+        repo_type="dataset",
+        local_dir=DATA_DIR,
+    )
+    size_mb = os.path.getsize(path) / 1024 / 1024
+    print(f"  {split}.jsonl: {size_mb:.1f} MB")
+
+print("Data downloaded!\n")
+
+# Step 2: Import and configure training
+import json
+import math
+import random
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+from huggingface_hub import HfApi
+import trackio
+
+
+# ============================================================
+# Configuration
+# ============================================================
+
+@dataclass
+class Config:
+    gene_model_id: str = "maayanlab/gsfm-rummagene"
+    text_model_id: str = "FremyCompany/BioLORD-2023"
+    shared_dim: int = 256
+    gene_dim: int = 256
+    text_dim: int = 768
+    proj_hidden_dim: int = 512
+    proj_dropout: float = 0.1
+    batch_size: int = 256
+    lr: float = 1e-4
+    gene_encoder_lr: float = 1e-5
+    weight_decay: float = 0.01
+    warmup_steps: int = 500
+    max_epochs: int = 50
+    patience: int = 10
+    temperature_init: float = 0.07
+    learnable_temperature: bool = True
+    gene_dropout_rate: float = 0.2
+    max_gene_set_size: int = 512
+    data_dir: str = DATA_DIR
+    output_dir: str = "/tmp/output"
+    hub_model_id: str = "AliSaadatV/GeneSetCLIP"
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    num_workers: int = 4
+    mixed_precision: bool = True
+    log_every: int = 10
+    eval_every: int = 1
+    save_every: int = 5
+
+
+# ============================================================
+# Dataset
+# ============================================================
+
+class GeneSetTextDataset(Dataset):
+    def __init__(self, jsonl_path, vocab, max_genes=512, gene_dropout=0.0, pad_idx=1):
+        self.records = []
+        with open(jsonl_path) as f:
+            for line in f:
+                self.records.append(json.loads(line))
+        self.vocab = vocab
+        self.max_genes = max_genes
+        self.gene_dropout = gene_dropout
+        self.pad_idx = pad_idx
+
+    def __len__(self):
+        return len(self.records)
+
+    def __getitem__(self, idx):
+        record = self.records[idx]
+        text = record["text"]
+        genes = record["genes"]
+        token_ids = [self.vocab.get(g, 0) for g in genes]
+
+        if self.gene_dropout > 0:
+            n_keep = max(3, int(len(token_ids) * (1 - self.gene_dropout)))
+            if n_keep < len(token_ids):
+                token_ids = random.sample(token_ids, n_keep)
+
+        if len(token_ids) > self.max_genes:
+            token_ids = random.sample(token_ids, self.max_genes)
+
+        n_genes = len(token_ids)
+        if n_genes < self.max_genes:
+            token_ids = token_ids + [self.pad_idx] * (self.max_genes - n_genes)
+
+        return {
+            "text": text,
+            "gene_ids": torch.tensor(token_ids, dtype=torch.long),
+            "n_genes": n_genes,
+            "id": record["id"],
+        }
+
+
+def collate_fn(batch):
+    return {
+        "text": [item["text"] for item in batch],
+        "gene_ids": torch.stack([item["gene_ids"] for item in batch]),
+        "n_genes": torch.tensor([item["n_genes"] for item in batch]),
+        "ids": [item["id"] for item in batch],
+    }
+
+
+# ============================================================
+# Model
+# ============================================================
+
+class ProjectionHead(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim, dropout=0.1):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, output_dim),
+            nn.LayerNorm(output_dim),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class GeneSetCLIP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.log_temperature = nn.Parameter(
+            torch.log(torch.tensor(config.temperature_init)),
+            requires_grad=config.learnable_temperature,
+        )
+        self.text_proj = ProjectionHead(config.text_dim, config.proj_hidden_dim,
+                                        config.shared_dim, config.proj_dropout)
+        self.gene_proj = ProjectionHead(config.gene_dim, config.shared_dim,
+                                        config.shared_dim, config.proj_dropout)
+
+    @property
+    def temperature(self):
+        return torch.clamp(self.log_temperature.exp(), min=0.01, max=1.0)
+
+    def forward(self, gene_emb, text_emb):
+        z_gene = F.normalize(self.gene_proj(gene_emb), dim=-1)
+        z_text = F.normalize(self.text_proj(text_emb), dim=-1)
+        tau = self.temperature
+        logits = z_gene @ z_text.T / tau
+        B = logits.size(0)
+        labels = torch.arange(B, device=logits.device)
+        loss_g2t = F.cross_entropy(logits, labels)
+        loss_t2g = F.cross_entropy(logits.T, labels)
+        loss = (loss_g2t + loss_t2g) / 2
+        with torch.no_grad():
+            g2t_acc = (logits.argmax(dim=1) == labels).float().mean()
+            t2g_acc = (logits.T.argmax(dim=1) == labels).float().mean()
+            avg_acc = (g2t_acc + t2g_acc) / 2
+        metrics = {
+            "loss": loss.item(), "g2t_acc": g2t_acc.item(),
+            "t2g_acc": t2g_acc.item(), "avg_acc": avg_acc.item(),
+            "temperature": tau.item(),
+        }
+        return loss, z_gene, z_text, metrics
+
+    def get_embeddings(self, gene_emb=None, text_emb=None):
+        z_gene = z_text = None
+        if gene_emb is not None:
+            z_gene = F.normalize(self.gene_proj(gene_emb), dim=-1)
+        if text_emb is not None:
+            z_text = F.normalize(self.text_proj(text_emb), dim=-1)
+        return z_gene, z_text
+
+
+# ============================================================
+# Evaluation
+# ============================================================
+
+@torch.no_grad()
+def evaluate_retrieval(model, gene_encoder, text_encoder, dataloader, device):
+    model.eval()
+    gene_encoder.eval()
+    all_z_gene, all_z_text, all_ids = [], [], []
+    total_loss, n_batches = 0, 0
+
+    for batch in dataloader:
+        gene_ids = batch["gene_ids"].to(device)
+        texts = batch["text"]
+        gene_emb = gene_encoder.encode(gene_ids)
+        text_emb = text_encoder.encode(texts, convert_to_tensor=True, show_progress_bar=False)
+        if text_emb.device != device:
+            text_emb = text_emb.to(device)
+        text_emb = text_emb.clone()
+        loss, z_gene, z_text, _ = model(gene_emb, text_emb)
+        total_loss += loss.item()
+        n_batches += 1
+        all_z_gene.append(z_gene.cpu())
+        all_z_text.append(z_text.cpu())
+        all_ids.extend(batch["ids"])
+
+    all_z_gene = torch.cat(all_z_gene, dim=0)
+    all_z_text = torch.cat(all_z_text, dim=0)
+    N = len(all_z_gene)
+    sim = all_z_gene @ all_z_text.T
+    labels = torch.arange(N)
+
+    def recall_at_k(sim_matrix, labels, k):
+        topk = sim_matrix.topk(min(k, sim_matrix.size(1)), dim=1).indices
+        return (topk == labels.unsqueeze(1)).any(dim=1).float().mean().item()
+
+    def mrr(sim_matrix, labels):
+        ranks = (sim_matrix.argsort(dim=1, descending=True) == labels.unsqueeze(1)).nonzero()[:, 1] + 1
+        return (1.0 / ranks.float()).mean().item()
+
+    results = {
+        "loss": total_loss / max(n_batches, 1), "n_samples": N,
+        "g2t_R@1": recall_at_k(sim, labels, 1),
+        "g2t_R@5": recall_at_k(sim, labels, 5),
+        "g2t_R@10": recall_at_k(sim, labels, 10),
+        "g2t_MRR": mrr(sim, labels),
+        "t2g_R@1": recall_at_k(sim.T, labels, 1),
+        "t2g_R@5": recall_at_k(sim.T, labels, 5),
+        "t2g_R@10": recall_at_k(sim.T, labels, 10),
+        "t2g_MRR": mrr(sim.T, labels),
+    }
+    results["avg_R@1"] = (results["g2t_R@1"] + results["t2g_R@1"]) / 2
+    results["avg_R@5"] = (results["g2t_R@5"] + results["t2g_R@5"]) / 2
+    results["avg_R@10"] = (results["g2t_R@10"] + results["t2g_R@10"]) / 2
+    results["avg_MRR"] = (results["g2t_MRR"] + results["t2g_MRR"]) / 2
+
+    model.train()
+    return results
+
+
+# ============================================================
+# Training
+# ============================================================
+
+def train(config):
+    print("=" * 70)
+    print("GeneSetCLIP Training")
+    print("=" * 70)
+    print(f"Device: {config.device}")
+    print(f"Batch size: {config.batch_size}")
+    print(f"Max epochs: {config.max_epochs}")
+
+    os.makedirs(config.output_dir, exist_ok=True)
+
+    # Load GSFM
+    print("\nLoading GSFM gene encoder...")
+    from gsfm import GSFM, Vocab
+    vocab_obj = Vocab.from_pretrained(config.gene_model_id)
+    gene_encoder = GSFM.from_pretrained(config.gene_model_id)
+    gene_encoder.to(config.device)
+    gene_encoder.train()
+    vocab_dict = {token: i for i, token in enumerate(vocab_obj.vocab)}
+    print(f"  GSFM vocab: {len(vocab_dict)} genes")
+
+    # Load BioLORD (frozen)
+    print("Loading BioLORD text encoder (frozen)...")
+    from sentence_transformers import SentenceTransformer
+    text_encoder = SentenceTransformer(config.text_model_id, device=config.device)
+    for param in text_encoder.parameters():
+        param.requires_grad = False
+    text_encoder.eval()
+
+    # Model
+    print("Building GeneSetCLIP model...")
+    model = GeneSetCLIP(config).to(config.device)
+    print(f"  Projection params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
+    print(f"  Gene encoder params: {sum(p.numel() for p in gene_encoder.parameters()):,}")
+
+    # Data
+    print("\nLoading datasets...")
+    train_ds = GeneSetTextDataset(os.path.join(config.data_dir, "train.jsonl"),
+                                   vocab_dict, config.max_gene_set_size, config.gene_dropout_rate)
+    val_ds = GeneSetTextDataset(os.path.join(config.data_dir, "val.jsonl"),
+                                 vocab_dict, config.max_gene_set_size, 0.0)
+    test_ds = GeneSetTextDataset(os.path.join(config.data_dir, "test.jsonl"),
+                                  vocab_dict, config.max_gene_set_size, 0.0)
+    print(f"  Train: {len(train_ds)}, Val: {len(val_ds)}, Test: {len(test_ds)}")
+
+    train_loader = DataLoader(train_ds, batch_size=config.batch_size, shuffle=True,
+                              collate_fn=collate_fn, num_workers=config.num_workers,
+                              pin_memory=True, drop_last=True)
+    val_loader = DataLoader(val_ds, batch_size=config.batch_size, shuffle=False,
+                            collate_fn=collate_fn, num_workers=config.num_workers)
+    test_loader = DataLoader(test_ds, batch_size=config.batch_size, shuffle=False,
+                             collate_fn=collate_fn, num_workers=config.num_workers)
+
+    steps_per_epoch = len(train_loader)
+    total_steps = steps_per_epoch * config.max_epochs
+    print(f"  Steps/epoch: {steps_per_epoch}, Total: {total_steps}")
+
+    # Optimizer
+    optimizer = torch.optim.AdamW([
+        {"params": list(model.text_proj.parameters()) + list(model.gene_proj.parameters()) +
+                   [model.log_temperature], "lr": config.lr, "weight_decay": config.weight_decay},
+        {"params": gene_encoder.parameters(), "lr": config.gene_encoder_lr,
+         "weight_decay": config.weight_decay},
+    ])
+
+    def lr_lambda(step):
+        if step < config.warmup_steps:
+            return step / max(config.warmup_steps, 1)
+        progress = (step - config.warmup_steps) / max(total_steps - config.warmup_steps, 1)
+        return 0.5 * (1 + math.cos(math.pi * progress))
+
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+    scaler = torch.amp.GradScaler('cuda') if config.mixed_precision and config.device == "cuda" else None
+
+    # Tracking
+    trackio.init(project="GeneSetCLIP",
+                 name=f"bs{config.batch_size}_lr{config.lr}_temp{config.temperature_init}")
+
+    # Training loop
+    best_val_mrr = 0
+    patience_counter = 0
+    global_step = 0
+
+    for epoch in range(1, config.max_epochs + 1):
+        model.train()
+        gene_encoder.train()
+        epoch_loss, epoch_acc, n_batches = 0, 0, 0
+
+        for batch in train_loader:
+            gene_ids = batch["gene_ids"].to(config.device)
+            texts = batch["text"]
+
+            gene_emb = gene_encoder.encode(gene_ids)
+            with torch.no_grad():
+                text_emb = text_encoder.encode(texts, convert_to_tensor=True, show_progress_bar=False)
+                if text_emb.device != torch.device(config.device):
+                    text_emb = text_emb.to(config.device)
+                text_emb = text_emb.clone()
+
+            if scaler is not None:
+                with torch.amp.autocast('cuda'):
+                    loss, _, _, metrics = model(gene_emb, text_emb)
+                optimizer.zero_grad()
+                scaler.scale(loss).backward()
+                scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(
+                    list(model.parameters()) + list(gene_encoder.parameters()), 1.0)
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                loss, _, _, metrics = model(gene_emb, text_emb)
+                optimizer.zero_grad()
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(
+                    list(model.parameters()) + list(gene_encoder.parameters()), 1.0)
+                optimizer.step()
+
+            scheduler.step()
+            global_step += 1
+            epoch_loss += metrics["loss"]
+            epoch_acc += metrics["avg_acc"]
+            n_batches += 1
+
+            if global_step % config.log_every == 0:
+                lr_p = optimizer.param_groups[0]["lr"]
+                lr_g = optimizer.param_groups[1]["lr"]
+                print(f"  Step {global_step:5d} | Loss: {metrics['loss']:.4f} | "
+                      f"Acc: {metrics['avg_acc']:.3f} | τ: {metrics['temperature']:.4f} | "
+                      f"LR: {lr_p:.2e}/{lr_g:.2e}")
+                trackio.log({
+                    "train/loss": metrics["loss"], "train/avg_acc": metrics["avg_acc"],
+                    "train/g2t_acc": metrics["g2t_acc"], "train/t2g_acc": metrics["t2g_acc"],
+                    "train/temperature": metrics["temperature"],
+                    "train/lr_proj": lr_p, "train/lr_gene": lr_g, "step": global_step,
+                })
+
+        avg_loss = epoch_loss / max(n_batches, 1)
+        avg_acc = epoch_acc / max(n_batches, 1)
+        print(f"\nEpoch {epoch}/{config.max_epochs} | Loss: {avg_loss:.4f} | Acc: {avg_acc:.3f}")
+
+        # Validation
+        if epoch % config.eval_every == 0:
+            print("  Evaluating...")
+            val_results = evaluate_retrieval(model, gene_encoder, text_encoder,
+                                             val_loader, config.device)
+            print(f"  Val R@1: {val_results['avg_R@1']:.3f} | R@5: {val_results['avg_R@5']:.3f} | "
+                  f"MRR: {val_results['avg_MRR']:.3f}")
+            trackio.log({
+                "val/loss": val_results["loss"],
+                "val/avg_R@1": val_results["avg_R@1"],
+                "val/avg_R@5": val_results["avg_R@5"],
+                "val/avg_R@10": val_results["avg_R@10"],
+                "val/avg_MRR": val_results["avg_MRR"],
+                "epoch": epoch,
+            })
+
+            if val_results["avg_MRR"] > best_val_mrr:
+                best_val_mrr = val_results["avg_MRR"]
+                patience_counter = 0
+                save_dir = os.path.join(config.output_dir, "best_model")
+                os.makedirs(save_dir, exist_ok=True)
+                torch.save(model.state_dict(), os.path.join(save_dir, "clip_model.pt"))
+                torch.save(gene_encoder.state_dict(), os.path.join(save_dir, "gene_encoder.pt"))
+                with open(os.path.join(save_dir, "config.json"), "w") as f:
+                    json.dump(vars(config), f, indent=2)
+                print(f"  ✓ New best! MRR: {best_val_mrr:.4f}")
+            else:
+                patience_counter += 1
+                print(f"  No improvement ({patience_counter}/{config.patience})")
+                if patience_counter >= config.patience:
+                    print(f"  Early stopping at epoch {epoch}")
+                    break
+
+    # Load best model
+    print("\n" + "=" * 70)
+    print("Final test evaluation...")
+    best_path = os.path.join(config.output_dir, "best_model")
+    if os.path.exists(best_path):
+        model.load_state_dict(torch.load(os.path.join(best_path, "clip_model.pt"),
+                                         map_location=config.device, weights_only=True))
+        gene_encoder.load_state_dict(torch.load(os.path.join(best_path, "gene_encoder.pt"),
+                                                map_location=config.device, weights_only=True))
+
+    test_results = evaluate_retrieval(model, gene_encoder, text_encoder,
+                                       test_loader, config.device)
+    print(f"Test Results:")
+    print(f"  G→T R@1: {test_results['g2t_R@1']:.3f}  R@5: {test_results['g2t_R@5']:.3f}  R@10: {test_results['g2t_R@10']:.3f}  MRR: {test_results['g2t_MRR']:.3f}")
+    print(f"  T→G R@1: {test_results['t2g_R@1']:.3f}  R@5: {test_results['t2g_R@5']:.3f}  R@10: {test_results['t2g_R@10']:.3f}  MRR: {test_results['t2g_MRR']:.3f}")
+    print(f"  Avg R@1: {test_results['avg_R@1']:.3f}  R@5: {test_results['avg_R@5']:.3f}  MRR: {test_results['avg_MRR']:.3f}")
+    trackio.log({"test/" + k: v for k, v in test_results.items()})
+
+    # Push to Hub
+    print("\nPushing to Hub...")
+    api = HfApi()
+    try:
+        api.create_repo(config.hub_model_id, exist_ok=True)
+    except Exception as e:
+        print(f"  Warning: {e}")
+
+    upload_dir = os.path.join(config.output_dir, "hub_upload")
+    os.makedirs(upload_dir, exist_ok=True)
+    torch.save(model.state_dict(), os.path.join(upload_dir, "clip_model.pt"))
+    torch.save(gene_encoder.state_dict(), os.path.join(upload_dir, "gene_encoder.pt"))
+    with open(os.path.join(upload_dir, "config.json"), "w") as f:
+        json.dump(vars(config), f, indent=2)
+    with open(os.path.join(upload_dir, "vocab.json"), "w") as f:
+        json.dump(vocab_dict, f)
+    with open(os.path.join(upload_dir, "test_results.json"), "w") as f:
+        json.dump(test_results, f, indent=2)
+
+    readme = f"""# GeneSetCLIP
+
+Contrastive model aligning gene-set embeddings (GSFM) with biomedical text descriptions (BioLORD-2023).
+
+## Architecture
+- **Gene encoder**: [GSFM](https://huggingface.co/maayanlab/gsfm-rummagene) (MLP autoencoder, 256-dim)
+- **Text encoder**: [BioLORD-2023](https://huggingface.co/FremyCompany/BioLORD-2023) (768-dim, frozen)
+- **Projection heads**: Maps both modalities to shared 256-dim space
+- **Loss**: Symmetric InfoNCE with learnable temperature
+
+## Training Data
+- **MSigDB v2024.1** (Human + Mouse): ~50,000 gene set-text pairs
+- Collections: H, C1-C8 (Human), MH, M1-M8 (Mouse)
+- Train: C2/C5/C8/C1 | Val: C3/C4 | Test: H/C6/C7
+
+## Test Results (H, C6, C7 — {test_results['n_samples']} gene sets)
+| Metric | Gene→Text | Text→Gene | Average |
+|--------|-----------|-----------|---------|
+| R@1    | {test_results['g2t_R@1']:.3f}     | {test_results['t2g_R@1']:.3f}     | {test_results['avg_R@1']:.3f}   |
+| R@5    | {test_results['g2t_R@5']:.3f}     | {test_results['t2g_R@5']:.3f}     | {test_results['avg_R@5']:.3f}   |
+| R@10   | {test_results['g2t_R@10']:.3f}    | {test_results['t2g_R@10']:.3f}    | {test_results['avg_R@10']:.3f}  |
+| MRR    | {test_results['g2t_MRR']:.3f}     | {test_results['t2g_MRR']:.3f}     | {test_results['avg_MRR']:.3f}   |
+
+## Usage
+
+```python
+import torch
+from gsfm import GSFM, Vocab
+from sentence_transformers import SentenceTransformer
+from huggingface_hub import hf_hub_download
+
+# Load gene encoder + vocab
+gene_encoder = GSFM.from_pretrained("maayanlab/gsfm-rummagene")
+vocab = Vocab.from_pretrained("maayanlab/gsfm-rummagene")
+gene_encoder.eval()
+
+# Load text encoder
+text_encoder = SentenceTransformer("FremyCompany/BioLORD-2023")
+
+# Load GeneSetCLIP projection heads
+clip_path = hf_hub_download("AliSaadatV/GeneSetCLIP", "clip_model.pt")
+config_path = hf_hub_download("AliSaadatV/GeneSetCLIP", "config.json")
+
+import json
+with open(config_path) as f:
+    cfg = json.load(f)
+
+# Reconstruct model (small — just projection heads)
+import torch.nn as nn, torch.nn.functional as F
+
+class ProjectionHead(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim, dropout=0.1):
+        super().__init__()
+        self.net = nn.Sequential(nn.Linear(input_dim, hidden_dim), nn.GELU(),
+                                 nn.Dropout(dropout), nn.Linear(hidden_dim, output_dim),
+                                 nn.LayerNorm(output_dim))
+    def forward(self, x): return self.net(x)
+
+class GeneSetCLIP(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.log_temperature = nn.Parameter(torch.zeros(1))
+        self.text_proj = ProjectionHead(768, 512, 256, 0.1)
+        self.gene_proj = ProjectionHead(256, 256, 256, 0.1)
+
+clip_model = GeneSetCLIP()
+clip_model.load_state_dict(torch.load(clip_path, map_location="cpu", weights_only=True))
+clip_model.eval()
+
+# Encode a gene set
+genes = ["TP53", "BRCA1", "EGFR", "MYC", "KRAS"]
+gene_ids = torch.tensor([vocab(genes)])
+with torch.no_grad():
+    gene_emb = gene_encoder.encode(gene_ids)
+    z_gene = F.normalize(clip_model.gene_proj(gene_emb), dim=-1)
+
+# Encode text
+text_emb = text_encoder.encode(["Tumor suppressor genes involved in cancer"],
+                                convert_to_tensor=True)
+with torch.no_grad():
+    z_text = F.normalize(clip_model.text_proj(text_emb), dim=-1)
+
+# Similarity
+print(f"Similarity: {{(z_gene @ z_text.T).item():.3f}}")
+```
+"""
+    with open(os.path.join(upload_dir, "README.md"), "w") as f:
+        f.write(readme)
+
+    api.upload_folder(folder_path=upload_dir, repo_id=config.hub_model_id,
+                      commit_message="Upload GeneSetCLIP trained model")
+    print(f"  Pushed to https://huggingface.co/{config.hub_model_id}")
+    print("\nDone!")
+
+
+if __name__ == "__main__":
+    config = Config()
+    train(config)