Add training script: CLIP-style multi-modal material embedding alignment

train_mattext_embeddings.py

@@ -0,0 +1,689 @@
+"""
+MatText Multi-Modal Embedding Alignment Training
+
+Architecture: CLIP-style contrastive learning across 8+ material text representations
+- Shared encoder (ModernBERT-base, 8192 ctx) with per-modality projection heads
+- All-pairs symmetric InfoNCE loss
+- Property-conditioned retrieval via property description encoding
+- FAISS vector database for cross-modal retrieval
+
+Based on:
+- MultiMat (AllPairsCLIP, arxiv:2312.00111)  
+- MatExpert (property↔structure InfoNCE, arxiv:2410.21317)
+- CrystalCLR (composition similarity, arxiv:2211.13408)
+
+Usage:
+    pip install torch transformers datasets faiss-cpu huggingface_hub trackio
+    python train_mattext_embeddings.py
+
+    # Or on HF Jobs:
+    # Hardware: a10g-large (24GB VRAM), timeout: 6h
+"""
+
+import os
+import json
+import math
+import time
+import logging
+import random
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from transformers import AutoModel, AutoTokenizer, get_cosine_schedule_with_warmup
+from datasets import load_dataset, concatenate_datasets
+from huggingface_hub import HfApi
+import faiss
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+
+# ============================================================================
+# Configuration
+# ============================================================================
+
+class Config:
+    # Model
+    encoder_name = "answerdotai/ModernBERT-base"
+    embed_dim = 128  # projection dimension (MultiMat recipe: 128-d)
+    max_length = 512  # tokens per modality input (ModernBERT supports up to 8192)
+    
+    # Modalities to align (columns in the dataset)
+    modalities = [
+        "composition",
+        "atom_sequences", 
+        "cif_symmetrized",
+        "cif_p1",
+        "zmatrix",
+        "atom_sequences_plusplus",
+        "slices",
+        "crystal_text_llm",
+        "local_env",
+        "robocrys_rep",  # natural language description (only in pretrain subsets)
+    ]
+    
+    # Training
+    batch_size = 32  
+    learning_rate = 2e-5
+    weight_decay = 0.01
+    num_epochs = 3
+    warmup_ratio = 0.1
+    temperature = 0.07  # InfoNCE temperature (MultiMat/CLIP standard)
+    grad_accum_steps = 8  # effective batch = 32*8 = 256 (critical for InfoNCE)
+    max_grad_norm = 1.0
+    gradient_checkpointing = True
+    max_modalities_per_step = 4  # randomly sample N modalities per step to save VRAM
+    
+    # Data
+    dataset_name = "n0w0f/MatText"
+    pretrain_config = "pretrain100k_v2"
+    finetune_configs = [
+        ("bandgap-train-filtered", "fold_0"),
+        ("form_energy-train-filtered", "fold_0"),
+    ]
+    max_train_samples = 50000
+    
+    # Output
+    output_dir = "mattext-embeddings"
+    hub_model_id = "n0w0f/mattext-aligned-embeddings"
+    push_to_hub = True
+    
+    # Device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    fp16 = torch.cuda.is_available()
+
+
+# ============================================================================
+# Model: Shared Encoder + Per-Modality Projection Heads
+# ============================================================================
+
+class ModalityProjection(nn.Module):
+    """2-layer MLP projection head (MultiMat recipe)"""
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, input_dim),
+            nn.GELU(),
+            nn.LayerNorm(input_dim),
+            nn.Linear(input_dim, output_dim),
+        )
+    
+    def forward(self, x):
+        return F.normalize(self.net(x), dim=-1)
+
+
+class MatTextEncoder(nn.Module):
+    """
+    Shared transformer encoder with per-modality projection heads.
+    All modalities share the same backbone but project to a shared 
+    embedding space through modality-specific heads.
+    """
+    def __init__(self, config: Config):
+        super().__init__()
+        self.config = config
+        
+        # Shared backbone
+        self.backbone = AutoModel.from_pretrained(config.encoder_name)
+        hidden_size = self.backbone.config.hidden_size
+        
+        if config.gradient_checkpointing:
+            self.backbone.gradient_checkpointing_enable()
+        
+        # Per-modality projection heads
+        self.projections = nn.ModuleDict({
+            mod: ModalityProjection(hidden_size, config.embed_dim)
+            for mod in config.modalities
+        })
+        
+        # Property projection (for property-conditioned queries)
+        self.property_projection = ModalityProjection(hidden_size, config.embed_dim)
+        
+        # Learnable temperature
+        self.log_temperature = nn.Parameter(
+            torch.tensor(math.log(1.0 / config.temperature))
+        )
+    
+    def encode(self, input_ids, attention_mask, modality_name):
+        """Encode a single modality"""
+        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
+        
+        # Mean pooling
+        mask = attention_mask.unsqueeze(-1).float()
+        hidden = outputs.last_hidden_state
+        pooled = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
+        
+        # Project through modality-specific head
+        if modality_name == "property":
+            return self.property_projection(pooled)
+        return self.projections[modality_name](pooled)
+    
+    @property
+    def temperature(self):
+        return torch.exp(self.log_temperature).clamp(min=0.01, max=100.0)
+
+
+# ============================================================================
+# Loss Functions
+# ============================================================================
+
+def symmetric_clip_loss(emb_a, emb_b, temperature):
+    """Symmetric InfoNCE (CLIP loss)"""
+    N = emb_a.size(0)
+    logits = (emb_a @ emb_b.T) * temperature
+    labels = torch.arange(N, device=emb_a.device)
+    loss_a = F.cross_entropy(logits, labels)
+    loss_b = F.cross_entropy(logits.T, labels)
+    return (loss_a + loss_b) / 2
+
+
+def all_pairs_clip_loss(embeddings_dict, temperature):
+    """AllPairsCLIP: sum symmetric InfoNCE over all modality pairs."""
+    mods = [k for k, v in embeddings_dict.items() if v is not None]
+    if len(mods) < 2:
+        return torch.tensor(0.0, requires_grad=True)
+    
+    device = embeddings_dict[mods[0]].device
+    total_loss = torch.tensor(0.0, device=device)
+    n_pairs = 0
+    
+    for i in range(len(mods)):
+        for j in range(i + 1, len(mods)):
+            total_loss = total_loss + symmetric_clip_loss(
+                embeddings_dict[mods[i]], embeddings_dict[mods[j]], temperature
+            )
+            n_pairs += 1
+    
+    return total_loss / n_pairs
+
+
+def property_similarity_loss(embeddings, labels, temperature):
+    """Property-aware soft contrastive loss (SupReMix-inspired)."""
+    N = embeddings.size(0)
+    if N < 2:
+        return torch.tensor(0.0, requires_grad=True)
+    
+    label_diff = torch.abs(labels.unsqueeze(0) - labels.unsqueeze(1))
+    max_diff = label_diff.max().clamp(min=1e-6)
+    label_sim = 1.0 - (label_diff / max_diff)
+    
+    cos_sim = embeddings @ embeddings.T
+    mask = torch.eye(N, device=embeddings.device).bool()
+    cos_sim = cos_sim.masked_fill(mask, 0)
+    label_sim = label_sim.masked_fill(mask, 0)
+    
+    return F.mse_loss(cos_sim, label_sim)
+
+
+# ============================================================================
+# Dataset
+# ============================================================================
+
+class MatTextMultiModalDataset(Dataset):
+    def __init__(self, data, modalities, property_col=None, property_name=None):
+        self.data = data
+        self.modalities = modalities
+        self.property_col = property_col
+        self.property_name = property_name
+        
+        available_cols = set(data.column_names) if hasattr(data, 'column_names') else set(data[0].keys())
+        self.available_modalities = [m for m in modalities if m in available_cols]
+        logger.info(f"Available modalities: {self.available_modalities}")
+        
+        self.has_properties = property_col is not None and property_col in available_cols
+        if self.has_properties:
+            logger.info(f"Property column '{property_col}' found")
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        row = self.data[idx]
+        item = {}
+        for mod in self.available_modalities:
+            text = row.get(mod, None)
+            if text and isinstance(text, str) and len(text.strip()) > 0:
+                item[mod] = text.strip()
+            else:
+                item[mod] = None
+        
+        if self.has_properties and row.get(self.property_col) is not None:
+            label_val = float(row[self.property_col])
+            comp = row.get("composition", "unknown")
+            item["property_text"] = f"composition: {comp} | {self.property_name}: {label_val:.4f}"
+            item["property_label"] = label_val
+        else:
+            item["property_text"] = None
+            item["property_label"] = None
+        
+        return item
+
+
+def collate_fn(batch, tokenizer, modalities, max_length):
+    result = {}
+    all_mod_keys = list(modalities) + ["property_text"]
+    
+    for mod in all_mod_keys:
+        texts = [item.get(mod) for item in batch]
+        valid_texts = [t for t in texts if t is not None]
+        if len(valid_texts) == 0:
+            result[mod] = None
+            continue
+        
+        texts_clean = [t if t is not None else "" for t in texts]
+        mask_valid = [t is not None for t in texts]
+        
+        encoded = tokenizer(texts_clean, padding=True, truncation=True, max_length=max_length, return_tensors="pt")
+        result[mod] = {
+            "input_ids": encoded["input_ids"],
+            "attention_mask": encoded["attention_mask"],
+            "valid_mask": torch.tensor(mask_valid, dtype=torch.bool),
+        }
+    
+    labels = [item.get("property_label") for item in batch]
+    if any(l is not None for l in labels):
+        labels_clean = [l if l is not None else 0.0 for l in labels]
+        labels_mask = [l is not None for l in labels]
+        result["property_labels"] = torch.tensor(labels_clean, dtype=torch.float32)
+        result["property_labels_mask"] = torch.tensor(labels_mask, dtype=torch.bool)
+    else:
+        result["property_labels"] = None
+        result["property_labels_mask"] = None
+    
+    return result
+
+
+# ============================================================================
+# Training Loop
+# ============================================================================
+
+def train_epoch(model, dataloader, optimizer, scheduler, config, epoch, scaler=None):
+    model.train()
+    total_loss = 0; total_clip_loss = 0; total_prop_loss = 0
+    log_interval = 20
+    
+    optimizer.zero_grad()
+    
+    for batch_idx, batch in enumerate(dataloader):
+        # Randomly sample modalities to save VRAM
+        available_mods = [m for m in config.modalities if batch.get(m) is not None]
+        if len(available_mods) > config.max_modalities_per_step:
+            must_have = [m for m in ["composition", "crystal_text_llm"] if m in available_mods]
+            remaining = [m for m in available_mods if m not in must_have]
+            n_sample = max(config.max_modalities_per_step - len(must_have), 1)
+            sampled = must_have + random.sample(remaining, min(n_sample, len(remaining)))
+        else:
+            sampled = available_mods
+        
+        embeddings = {}
+        for mod in sampled:
+            if batch.get(mod) is None:
+                embeddings[mod] = None; continue
+            
+            input_ids = batch[mod]["input_ids"].to(config.device)
+            attention_mask = batch[mod]["attention_mask"].to(config.device)
+            valid_mask = batch[mod]["valid_mask"]
+            
+            if not valid_mask.any():
+                embeddings[mod] = None; continue
+            
+            with torch.amp.autocast('cuda', enabled=config.fp16):
+                emb = model.encode(input_ids, attention_mask, mod)
+            emb = emb * valid_mask.to(config.device).unsqueeze(-1).float()
+            embeddings[mod] = emb
+        
+        with torch.amp.autocast('cuda', enabled=config.fp16):
+            temperature = model.temperature
+            clip_l = all_pairs_clip_loss(embeddings, temperature)
+        
+        prop_l = torch.tensor(0.0, device=config.device)
+        if batch.get("property_text") is not None and batch.get("property_labels") is not None:
+            prop_ids = batch["property_text"]["input_ids"].to(config.device)
+            prop_mask = batch["property_text"]["attention_mask"].to(config.device)
+            prop_valid = batch["property_text"]["valid_mask"]
+            
+            if prop_valid.any():
+                with torch.amp.autocast('cuda', enabled=config.fp16):
+                    prop_emb = model.encode(prop_ids, prop_mask, "property")
+                
+                labels = batch["property_labels"].to(config.device)
+                labels_mask = batch["property_labels_mask"].to(config.device)
+                
+                if labels_mask.sum() > 1:
+                    prop_l = property_similarity_loss(prop_emb[labels_mask], labels[labels_mask], temperature)
+                
+                for anchor_mod in ["robocrys_rep", "crystal_text_llm", "composition"]:
+                    if embeddings.get(anchor_mod) is not None:
+                        with torch.amp.autocast('cuda', enabled=config.fp16):
+                            prop_clip = symmetric_clip_loss(
+                                prop_emb[labels_mask], embeddings[anchor_mod][labels_mask], temperature
+                            )
+                        prop_l = prop_l + 0.5 * prop_clip
+                        break
+        
+        loss = (clip_l + 0.3 * prop_l) / config.grad_accum_steps
+        
+        if config.fp16 and scaler is not None:
+            scaler.scale(loss).backward()
+        else:
+            loss.backward()
+        
+        if (batch_idx + 1) % config.grad_accum_steps == 0:
+            if config.fp16 and scaler is not None:
+                scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm)
+                scaler.step(optimizer); scaler.update()
+            else:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm)
+                optimizer.step()
+            scheduler.step(); optimizer.zero_grad()
+        
+        total_loss += loss.item() * config.grad_accum_steps
+        total_clip_loss += clip_l.item()
+        total_prop_loss += prop_l.item() if isinstance(prop_l, torch.Tensor) else prop_l
+        
+        if (batch_idx + 1) % log_interval == 0:
+            avg = total_loss / (batch_idx + 1)
+            logger.info(
+                f"Epoch {epoch} | {batch_idx+1}/{len(dataloader)} | "
+                f"Loss: {avg:.4f} | CLIP: {total_clip_loss/(batch_idx+1):.4f} | "
+                f"Prop: {total_prop_loss/(batch_idx+1):.4f} | "
+                f"LR: {scheduler.get_last_lr()[0]:.2e} | T: {model.temperature.item():.3f}"
+            )
+    
+    return total_loss / max(len(dataloader), 1)
+
+
+# ============================================================================
+# Evaluation
+# ============================================================================
+
+@torch.no_grad()
+def evaluate_retrieval(model, dataloader, config, k_values=[1, 5, 10]):
+    model.eval()
+    all_embeddings = {mod: [] for mod in config.modalities}
+    
+    for batch in dataloader:
+        for mod in config.modalities:
+            if batch.get(mod) is None: continue
+            input_ids = batch[mod]["input_ids"].to(config.device)
+            attention_mask = batch[mod]["attention_mask"].to(config.device)
+            valid_mask = batch[mod]["valid_mask"]
+            if not valid_mask.any(): continue
+            
+            emb = model.encode(input_ids, attention_mask, mod).cpu()
+            for i in range(len(emb)):
+                all_embeddings[mod].append(emb[i] if valid_mask[i] else None)
+    
+    results = {}
+    eval_pairs = [
+        ("composition", "crystal_text_llm"), ("composition", "cif_symmetrized"),
+        ("slices", "crystal_text_llm"), ("composition", "slices"),
+    ]
+    if len([e for e in all_embeddings.get("robocrys_rep", []) if e is not None]) > 0:
+        eval_pairs.extend([("robocrys_rep", "composition"), ("robocrys_rep", "cif_symmetrized")])
+    
+    for mod_a, mod_b in eval_pairs:
+        embs_a, embs_b = all_embeddings.get(mod_a, []), all_embeddings.get(mod_b, [])
+        if not embs_a or not embs_b: continue
+        
+        valid_idx = [i for i in range(min(len(embs_a), len(embs_b))) 
+                     if embs_a[i] is not None and embs_b[i] is not None]
+        if len(valid_idx) < 10: continue
+        
+        ea = torch.stack([embs_a[i] for i in valid_idx])
+        eb = torch.stack([embs_b[i] for i in valid_idx])
+        sim = ea @ eb.T
+        
+        recalls = {}
+        for k in k_values:
+            kk = min(k, len(valid_idx) - 1)
+            topk = sim.topk(kk, dim=1).indices
+            correct = (topk == torch.arange(len(valid_idx)).unsqueeze(1)).any(dim=1)
+            recalls[f"R@{k}"] = correct.float().mean().item()
+        
+        results[f"{mod_a}→{mod_b}"] = recalls
+        logger.info(f"  {mod_a}→{mod_b}: {recalls}")
+    
+    return results
+
+
+# ============================================================================
+# FAISS Vector Database
+# ============================================================================
+
+def build_vector_database(model, dataset, tokenizer, config, modalities_to_index=None):
+    if modalities_to_index is None:
+        modalities_to_index = config.modalities
+    model.eval()
+    
+    all_embeddings = {mod: [] for mod in modalities_to_index}
+    all_metadata = []
+    bs = 64
+    
+    for start in range(0, len(dataset), bs):
+        end = min(start + bs, len(dataset))
+        items = [dataset[i] for i in range(start, end)]
+        batch = collate_fn(items, tokenizer, config.modalities, config.max_length)
+        
+        for item in items:
+            all_metadata.append({"composition": item.get("composition", ""), "property_label": item.get("property_label")})
+        
+        with torch.no_grad():
+            for mod in modalities_to_index:
+                if batch.get(mod) is None:
+                    all_embeddings[mod].extend([None] * len(items)); continue
+                emb = model.encode(
+                    batch[mod]["input_ids"].to(config.device),
+                    batch[mod]["attention_mask"].to(config.device), mod
+                ).cpu().numpy()
+                for i in range(len(emb)):
+                    all_embeddings[mod].append(emb[i] if batch[mod]["valid_mask"][i] else None)
+        
+        if (start // bs) % 10 == 0:
+            logger.info(f"Indexed {end}/{len(dataset)}")
+    
+    indices = {}
+    for mod in modalities_to_index:
+        valid_embs = [e for e in all_embeddings[mod] if e is not None]
+        valid_map = [i for i, e in enumerate(all_embeddings[mod]) if e is not None]
+        if not valid_embs: continue
+        
+        emb_matrix = np.stack(valid_embs).astype(np.float32)
+        faiss.normalize_L2(emb_matrix)
+        d = emb_matrix.shape[1]
+        
+        if len(valid_embs) > 10000:
+            nlist = min(100, int(np.sqrt(len(valid_embs))))
+            q = faiss.IndexFlatIP(d)
+            index = faiss.IndexIVFFlat(q, d, nlist, faiss.METRIC_INNER_PRODUCT)
+            index.train(emb_matrix)
+        else:
+            index = faiss.IndexFlatIP(d)
+        index.add(emb_matrix)
+        
+        indices[mod] = {"index": index, "valid_indices_map": valid_map,
+                        "metadata": [all_metadata[i] for i in valid_map]}
+        logger.info(f"FAISS {mod}: {len(valid_embs)} vectors, dim={d}")
+    
+    return indices
+
+
+def search_vector_db(query_text, query_modality, model, tokenizer, indices, config, k=10):
+    model.eval()
+    enc = tokenizer([query_text], padding=True, truncation=True, max_length=config.max_length, return_tensors="pt")
+    with torch.no_grad():
+        q = model.encode(enc["input_ids"].to(config.device), enc["attention_mask"].to(config.device), query_modality)
+    q = q.cpu().numpy().astype(np.float32)
+    faiss.normalize_L2(q)
+    
+    results = []
+    for mod_name, idx_data in indices.items():
+        scores, ids = idx_data["index"].search(q, k)
+        for s, i in zip(scores[0], ids[0]):
+            if i >= 0:
+                m = dict(idx_data["metadata"][i])
+                m["matched_modality"] = mod_name
+                results.append((float(s), m))
+    
+    results.sort(key=lambda x: x[0], reverse=True)
+    seen, unique = set(), []
+    for s, m in results:
+        c = m.get("composition", "")
+        if c not in seen:
+            seen.add(c); unique.append((s, m))
+            if len(unique) >= k: break
+    return unique
+
+
+# ============================================================================
+# Main
+# ============================================================================
+
+def main():
+    config = Config()
+    logger.info(f"Device: {config.device} | Encoder: {config.encoder_name}")
+    logger.info(f"Batch: {config.batch_size}x{config.grad_accum_steps}={config.batch_size*config.grad_accum_steps}")
+    
+    try:
+        import trackio
+        trackio.init(project="mattext-embeddings", name=f"align-{config.encoder_name.split('/')[-1]}")
+        use_trackio = True
+    except:
+        use_trackio = False
+    
+    tokenizer = AutoTokenizer.from_pretrained(config.encoder_name)
+    model = MatTextEncoder(config).to(config.device)
+    logger.info(f"Params: {sum(p.numel() for p in model.parameters()):,}")
+    
+    # Load data
+    pretrain_data = load_dataset(config.dataset_name, config.pretrain_config, split="train")
+    logger.info(f"Pretrain: {len(pretrain_data)} samples, cols: {pretrain_data.column_names}")
+    
+    finetune_data = None
+    for ft_cfg, ft_split in config.finetune_configs:
+        try:
+            ft = load_dataset(config.dataset_name, ft_cfg, split=ft_split)
+            logger.info(f"Loaded {ft_cfg}/{ft_split}: {len(ft)} samples")
+            finetune_data = ft if finetune_data is None else concatenate_datasets([
+                finetune_data.select_columns(list(set(finetune_data.column_names) & set(ft.column_names))),
+                ft.select_columns(list(set(finetune_data.column_names) & set(ft.column_names)))
+            ])
+        except Exception as e:
+            logger.warning(f"Failed {ft_cfg}: {e}")
+    
+    if len(pretrain_data) > config.max_train_samples:
+        pretrain_data = pretrain_data.shuffle(seed=42).select(range(config.max_train_samples))
+    
+    make_collate = lambda tok, mods, ml: lambda batch: collate_fn(batch, tok, mods, ml)
+    
+    pretrain_loader = DataLoader(
+        MatTextMultiModalDataset(pretrain_data, config.modalities),
+        batch_size=config.batch_size, shuffle=True, drop_last=True, num_workers=0,
+        collate_fn=make_collate(tokenizer, config.modalities, config.max_length),
+        pin_memory=config.device == "cuda",
+    )
+    
+    finetune_loader = None
+    if finetune_data:
+        if len(finetune_data) > config.max_train_samples:
+            finetune_data = finetune_data.shuffle(seed=42).select(range(config.max_train_samples))
+        finetune_loader = DataLoader(
+            MatTextMultiModalDataset(finetune_data, config.modalities, "labels", "property_value"),
+            batch_size=config.batch_size, shuffle=True, drop_last=True, num_workers=0,
+            collate_fn=make_collate(tokenizer, config.modalities, config.max_length),
+            pin_memory=config.device == "cuda",
+        )
+    
+    optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay)
+    total_steps = len(pretrain_loader) * config.num_epochs // config.grad_accum_steps
+    if finetune_loader:
+        total_steps += len(finetune_loader) * config.num_epochs // config.grad_accum_steps
+    scheduler = get_cosine_schedule_with_warmup(optimizer, int(total_steps * config.warmup_ratio), total_steps)
+    scaler = torch.amp.GradScaler('cuda') if config.fp16 else None
+    
+    logger.info(f"Steps: {total_steps}")
+    
+    # Phase 1: Multi-modal alignment
+    logger.info("=" * 60 + "\nPhase 1: Multi-modal alignment\n" + "=" * 60)
+    best_loss = float('inf')
+    for epoch in range(1, config.num_epochs + 1):
+        t0 = time.time()
+        loss = train_epoch(model, pretrain_loader, optimizer, scheduler, config, epoch, scaler)
+        logger.info(f"Epoch {epoch} | Loss: {loss:.4f} | Time: {time.time()-t0:.0f}s")
+        if use_trackio:
+            try: trackio.log({"phase": 1, "epoch": epoch, "loss": loss})
+            except: pass
+        if loss < best_loss:
+            best_loss = loss
+            os.makedirs(config.output_dir, exist_ok=True)
+            torch.save(model.state_dict(), f"{config.output_dir}/best_model.pt")
+    
+    # Phase 2: Property-conditioned alignment
+    if finetune_loader:
+        logger.info("=" * 60 + "\nPhase 2: Property-conditioned alignment\n" + "=" * 60)
+        for epoch in range(1, config.num_epochs + 1):
+            t0 = time.time()
+            loss = train_epoch(model, finetune_loader, optimizer, scheduler, config, epoch, scaler)
+            logger.info(f"P2 Epoch {epoch} | Loss: {loss:.4f} | Time: {time.time()-t0:.0f}s")
+            if loss < best_loss:
+                best_loss = loss
+                torch.save(model.state_dict(), f"{config.output_dir}/best_model.pt")
+    
+    # Evaluate
+    logger.info("=" * 60 + "\nEvaluation\n" + "=" * 60)
+    eval_data = load_dataset(config.dataset_name, config.pretrain_config, split="test")
+    if len(eval_data) > 5000:
+        eval_data = eval_data.shuffle(seed=42).select(range(5000))
+    
+    eval_loader = DataLoader(
+        MatTextMultiModalDataset(eval_data, config.modalities),
+        batch_size=config.batch_size, shuffle=False, num_workers=0,
+        collate_fn=make_collate(tokenizer, config.modalities, config.max_length),
+    )
+    results = evaluate_retrieval(model, eval_loader, config)
+    
+    # Build FAISS DB
+    logger.info("Building FAISS indices...")
+    db = build_vector_database(
+        model, MatTextMultiModalDataset(eval_data, config.modalities),
+        tokenizer, config, ["composition", "crystal_text_llm", "slices", "cif_symmetrized"]
+    )
+    
+    os.makedirs(f"{config.output_dir}/faiss", exist_ok=True)
+    for mod, d in db.items():
+        faiss.write_index(d["index"], f"{config.output_dir}/faiss/{mod}.index")
+        with open(f"{config.output_dir}/faiss/{mod}_metadata.json", "w") as f:
+            json.dump(d["metadata"], f)
+    
+    # Demo
+    for q, m in [("Fe2O3", "composition"), ("Si Ge", "composition")]:
+        logger.info(f"\nQuery: '{q}' ({m})")
+        for rank, (s, meta) in enumerate(search_vector_db(q, m, model, tokenizer, db, config, 5), 1):
+            logger.info(f"  #{rank}: {s:.4f} | {meta}")
+    
+    # Save & push
+    torch.save(model.state_dict(), f"{config.output_dir}/model.pt")
+    tokenizer.save_pretrained(config.output_dir)
+    with open(f"{config.output_dir}/config.json", "w") as f:
+        json.dump({k: str(v) if not isinstance(v, (int, float, str, bool, list, dict, type(None))) else v
+                   for k, v in vars(Config).items() if not k.startswith("_")}, f, indent=2)
+    with open(f"{config.output_dir}/retrieval_results.json", "w") as f:
+        json.dump(results, f, indent=2)
+    
+    if config.push_to_hub:
+        try:
+            api = HfApi()
+            api.create_repo(config.hub_model_id, exist_ok=True)
+            api.upload_folder(folder_path=config.output_dir, repo_id=config.hub_model_id,
+                            commit_message="Upload MatText aligned embeddings + FAISS indices")
+            logger.info(f"Pushed to https://huggingface.co/{config.hub_model_id}")
+        except Exception as e:
+            logger.error(f"Push failed: {e}")
+    
+    logger.info("DONE!")
+
+
+if __name__ == "__main__":
+    main()