Initial commit

test_model.py

@@ -0,0 +1,96 @@
+import argparse
+import os
+import sys
+import numpy as np
+import pandas as pd
+import torch
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
+from sklearn.model_selection import cross_validate
+
+# Конфигурация по умолчанию
+DEFAULT_TASKS = ['ESOL', 'FreeSolv', 'HIV', 'BACE', 'BBBP', 'ClinTox']
+MODEL_NAME = "DeepChem/ChemBERTa-10M-MLM"
+
+def load_model_and_checkpoint(checkpoint_path, device="cpu"):
+    print(f"Loading model {MODEL_NAME}...", file=sys.stderr)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    backbone = AutoModelForMaskedLM.from_pretrained(MODEL_NAME).roberta.to(device)
+
+    if not os.path.exists(checkpoint_path):
+        raise FileNotFoundError(f"File not found: {checkpoint_path}")
+
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    backbone.load_state_dict(checkpoint['backbone'])
+    backbone.eval()
+    print("Model is loaded", file=sys.stderr)
+    return tokenizer, backbone
+
+@torch.no_grad()
+def mol_to_emb(smiles, tokenizer, model, device="cpu"):
+    tokenized = tokenizer([smiles], padding=False, return_tensors="pt")
+    input_ids = tokenized['input_ids'].to(device)
+    hs = model(input_ids).last_hidden_state
+
+    emb = torch.cat([hs[:, 0], hs[:, 1:].mean(dim=1)], dim=1)
+    return emb.squeeze(0).cpu().numpy()
+
+def evaluate_tasks(checkpoint_path, data_dir='./support/', device="cpu"):
+    tasks = DEFAULT_TASKS
+    tokenizer, model = load_model_and_checkpoint(checkpoint_path, device)
+
+    results = {}
+    for task in tasks:
+        csv_path = os.path.join(data_dir, f"{task}.csv")
+        if not os.path.exists(csv_path):
+            print(f"\n[WARN] File {csv_path} not found. Skipping '{task}'.", file=sys.stderr)
+            continue
+
+        print(f"Task: {task}", file=sys.stderr)
+        ds = pd.read_csv(csv_path, sep='\t')
+
+        # Вычисление эмбеддингов с прогресс-баром
+        ds['v'] = ds['X'].apply(lambda x: mol_to_emb(x, tokenizer, model, device))
+        ds = ds.sample(frac=1, random_state=42).reset_index(drop=True)
+
+        # Подготовка данных для sklearn
+        X = np.stack(ds['v'].values)
+        y = ds['y'].to_numpy()
+
+        # Выбор модели и метрики
+        if task in ['ESOL', 'FreeSolv']:
+            rf_model = RandomForestRegressor(random_state=42, n_jobs=5)
+            scoring = 'neg_mean_absolute_error'
+            metric_name = "MAE"
+        else:
+            rf_model = RandomForestClassifier(random_state=42, n_jobs=5)
+            scoring = 'f1_macro'
+            metric_name = "F1-macro"
+
+        # Кросс-валидация
+        cv_results = cross_validate(rf_model, X, y, cv=5, scoring=scoring, n_jobs=1)
+        mean_score = cv_results['test_score'].mean()
+        std_score = cv_results['test_score'].std()
+        results[task] = (mean_score, std_score)
+        print(f"  {metric_name}: {mean_score:.4f} ± {std_score:.4f}", file=sys.stderr)
+
+
+    for task, (mean, std) in results.items():
+        print(f"{task:10}: {mean:.4f} ± {std:.4f}")
+    return results
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "checkpoint_path", type=str,
+        help="Path to checkpoint file (.pth)"
+    )
+    parser.add_argument(
+        "--device", type=str, default="cpu", choices=["cpu", "cuda"],
+    )
+    args = parser.parse_args()
+
+    evaluate_tasks(
+        checkpoint_path=args.checkpoint_path,
+        device=args.device
+    )

train_model.py

@@ -0,0 +1,312 @@
+import os
+import time
+import math
+import logging
+import argparse
+from datetime import datetime
+from typing import Dict, List, Any
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer,
+    AutoModelForMaskedLM,
+    DataCollatorForLanguageModeling,
+    PreTrainedTokenizerBase
+)
+from rdkit import Chem
+from rdkit.Chem import Descriptors
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
+logger = logging.getLogger(__name__)
+
+def compute_rdkit_features(smiles: str) -> np.ndarray:
+    try:
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            return np.zeros(210, dtype=np.float32)
+        return np.array(list(Descriptors.CalcMolDescriptors(mol).values()))
+    except Exception:
+        return np.zeros(210, dtype=np.float32)
+
+class SMILESAndDescriptorCollator:
+    def __init__(
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        max_length: int = 512,
+        mlm_probability: float = 0.15,
+        do_mlm: bool = True
+    ):
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.do_mlm = do_mlm
+        if self.do_mlm:
+            self.mlm_collator = DataCollatorForLanguageModeling(
+                tokenizer=self.tokenizer,
+                mlm=True,
+                mlm_probability=mlm_probability,
+                return_tensors="pt"
+            )
+        else:
+            self.mlm_collator = None
+
+    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
+        smiles_batch = [f['smiles'] for f in features]
+        descriptors_list = [f['descriptors'] for f in features]
+
+        tokenized = self.tokenizer(
+            smiles_batch,
+            padding=False,
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors=None
+        )
+
+        features_for_mlm = [
+            {k: v[i] for k, v in tokenized.items()}
+            for i in range(len(smiles_batch))
+        ]
+
+        if self.do_mlm and self.mlm_collator:
+            batch_text = self.mlm_collator(features_for_mlm)
+        else:
+            tokenized_padded = self.tokenizer.pad(
+                features_for_mlm,
+                padding=True,
+                max_length=self.max_length,
+                return_tensors="pt"
+            )
+            batch_text = dict(tokenized_padded)
+
+        descriptors_tensor = torch.tensor(np.stack(descriptors_list), dtype=torch.float32)
+
+        batch = batch_text
+        batch['descriptors'] = descriptors_tensor
+
+        return batch
+
+def get_backbone_grad_vector(module, exclude_keywords=None):
+    if exclude_keywords is None:
+        exclude_keywords = []
+
+    grads = []
+    for name, param in module.named_parameters():
+        if any(keyword in name.lower() for keyword in exclude_keywords):
+            continue
+        if param.grad is not None:
+            grads.append(param.grad.detach().flatten())
+
+    if len(grads) == 0:
+        return torch.tensor([])
+
+    return torch.cat(grads)
+
+def compute_gradient_metrics(model, loss1, loss2, exclude_keywords=None):
+    if exclude_keywords is None:
+        exclude_keywords = []
+
+    model.zero_grad(set_to_none=True)
+    loss1.backward(retain_graph=True)
+    g1 = get_backbone_grad_vector(model, exclude_keywords)
+    norm_mtr = g1.norm().item() if g1 is not None and g1.numel() > 0 else None
+
+    model.zero_grad(set_to_none=True)
+    loss2.backward(retain_graph=True)
+    g2 = get_backbone_grad_vector(model, exclude_keywords)
+    norm_mlm = g2.norm().item() if g2 is not None and g2.numel() > 0 else None
+
+    model.zero_grad(set_to_none=True)
+
+    angle_deg = None
+    if (g1 is not None and g2 is not None and
+        g1.numel() > 0 and g2.numel() > 0 and
+        g1.numel() == g2.numel()):
+        cos_sim = F.cosine_similarity(g1.unsqueeze(0), g2.unsqueeze(0), dim=1).item()
+        cos_sim = max(min(cos_sim, 1.0), -1.0)
+        angle_rad = math.acos(cos_sim)
+        angle_deg = math.degrees(angle_rad)
+
+    return {
+        'angle_deg': angle_deg,
+        'norm_mtr': norm_mtr,
+        'norm_mlm': norm_mlm
+    }
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description="ChemBERTa Multi-Task Training")
+    parser.add_argument("--smiles_file", type=str, default="support/smiles_10k.txt")
+    parser.add_argument("--stats_file", type=str, default="support/normalization_params.pth")
+    parser.add_argument("--output_file", type=str, default="model.pth")
+    parser.add_argument("--batch_size", type=int, default=64)
+    parser.add_argument("--max_length", type=int, default=128)
+    parser.add_argument("--mlm_weight", type=float, default=1.0)
+    parser.add_argument("--mtr_weight", type=float, default=1.0)
+    parser.add_argument("--lr", type=float, default=3e-5)
+    parser.add_argument("--epochs", type=int, default=1)
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    logger.info(f"Using device: {device}")
+
+    logger.info("Loading model...")
+    tokenizer = AutoTokenizer.from_pretrained("DeepChem/ChemBERTa-10M-MLM")
+    model_base = AutoModelForMaskedLM.from_pretrained("DeepChem/ChemBERTa-10M-MLM").roberta
+    model_dim = 384
+
+    mlm_head = nn.Sequential(
+        nn.Linear(model_dim, model_dim * 2),
+        nn.GELU(),
+        nn.Linear(model_dim * 2, tokenizer.vocab_size),
+    )
+    rdkit_head = nn.Sequential(
+        nn.Linear(model_dim, model_dim * 2),
+        nn.GELU(),
+        nn.Linear(model_dim * 2, 210),
+    )
+
+    model_base.to(device)
+    mlm_head.to(device)
+    rdkit_head.to(device)
+
+    logger.info("Loading dataset...")
+    raw_dataset = load_dataset("text", data_files={"train": args.smiles_file})
+    raw_dataset = raw_dataset.rename_column("text", "smiles")
+
+    logger.info("Calculating RDKit features...")
+    processed_dataset = raw_dataset.map(
+        lambda x: {"descriptors": compute_rdkit_features(x["smiles"])},
+        num_proc=8,
+        desc="Calculating RDKit features"
+    )
+
+    collator = SMILESAndDescriptorCollator(tokenizer=tokenizer, max_length=args.max_length)
+    dataloader = DataLoader(
+        processed_dataset["train"],
+        batch_size=args.batch_size,
+        collate_fn=collator
+    )
+
+    logger.info("Loading normalization stats...")
+    stats = torch.load(args.stats_file, map_location=device)
+    means = stats["means"].to(device)
+    stds = stats["stds"].to(device)
+    stds[stds < 1e-6] = 1.0
+
+    optimizer = torch.optim.AdamW(
+        list(model_base.parameters()) + list(mlm_head.parameters()) + list(rdkit_head.parameters()),
+        lr=args.lr, weight_decay=1e-4
+    )
+
+    clip_grad_norm = 1.0
+    BACKBONE_EXCLUDE_KEYWORDS = ["head", "rdkit", "mlm", "classifier", "pooler"]
+    LOG_GRAD_METRICS_EVERY_N_BATCHES = 10
+
+    log_dir = os.path.join("runs", f"train_{datetime.now().strftime('%Y%m%d_%H%M%S')}")
+    writer = SummaryWriter(log_dir=log_dir)
+    global_step = 0
+
+    logger.info("Starting training")
+
+    for epoch in range(args.epochs):
+        start_time = time.time()
+        total_loss_mtr = 0.0
+        total_loss_mlm = 0.0
+        total_loss = 0.0
+        num_batches = 0
+
+        pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{args.epochs}")
+        model_base.train()
+        mlm_head.train()
+        rdkit_head.train()
+
+        for batch_idx, batch in enumerate(pbar):
+            input_ids = batch["input_ids"].to(device)
+            attention_mask = batch["attention_mask"].to(device)
+            descriptors = batch["descriptors"].to(device)
+            labels = batch["labels"].to(device)
+
+            outputs = model_base(input_ids, attention_mask=attention_mask)
+            result = outputs.last_hidden_state
+
+            mtr_res = rdkit_head(result[:, 0])
+            mlm_res = mlm_head(result)
+
+            loss_mtr = F.huber_loss(mtr_res, (descriptors - means) / stds)
+            loss_mlm = F.cross_entropy(mlm_res.flatten(end_dim=1), labels.flatten(), ignore_index=-100)
+            loss = loss_mtr * args.mtr_weight + loss_mlm * args.mlm_weight
+
+            if num_batches % LOG_GRAD_METRICS_EVERY_N_BATCHES == 0:
+                metrics = compute_gradient_metrics(
+                    model=model_base, loss1=loss_mtr, loss2=loss_mlm,
+                    exclude_keywords=BACKBONE_EXCLUDE_KEYWORDS
+                )
+                postfix = {
+                    "loss_mlm": f"{loss_mlm.item():.4f}",
+                    "loss_mtr": f"{loss_mtr.item():.4f}",
+                }
+                if metrics["angle_deg"] is not None:
+                    postfix["angle"] = f"{metrics['angle_deg']:.1f}°"
+                    writer.add_scalar("gradients/backbone_angle_deg", metrics["angle_deg"], global_step)
+                if metrics["norm_mtr"] is not None:
+                    postfix["‖∇MTR‖"] = f"{metrics['norm_mtr']:.3f}"
+                    writer.add_scalar("gradients/backbone_norm_mtr", metrics["norm_mtr"], global_step)
+                if metrics["norm_mlm"] is not None:
+                    postfix["‖∇MLM‖"] = f"{metrics['norm_mlm']:.3f}"
+                    writer.add_scalar("gradients/backbone_norm_mlm", metrics["norm_mlm"], global_step)
+
+                pbar.set_postfix(postfix)
+
+                writer.add_scalar("loss/total", loss.item(), global_step)
+                writer.add_scalar("loss/mtr_l1", loss_mtr.item(), global_step)
+                writer.add_scalar("loss/mlm_ce", loss_mlm.item(), global_step)
+
+            optimizer.zero_grad()
+            loss.backward()
+
+            grad_norm = torch.nn.utils.clip_grad_norm_(model_base.parameters(), clip_grad_norm)
+            torch.nn.utils.clip_grad_norm_(rdkit_head.parameters(), clip_grad_norm)
+            torch.nn.utils.clip_grad_norm_(mlm_head.parameters(), clip_grad_norm)
+
+            writer.add_scalar("training/grad_norm_clipped", grad_norm.item(), global_step)
+            writer.add_scalar("training/learning_rate", optimizer.param_groups[0]["lr"], global_step)
+
+            optimizer.step()
+
+            total_loss += loss.item()
+            total_loss_mtr += loss_mtr.item()
+            total_loss_mlm += loss_mlm.item()
+            num_batches += 1
+            global_step += 1
+
+        epoch_time = time.time() - start_time
+        avg_loss = total_loss / num_batches
+        avg_loss_mtr = total_loss_mtr / num_batches
+        avg_loss_mlm = total_loss_mlm / num_batches
+
+        writer.add_scalar("epoch/avg_total_loss", avg_loss, epoch)
+        writer.add_scalar("epoch/avg_loss_mtr", avg_loss_mtr, epoch)
+        writer.add_scalar("epoch/avg_loss_mlm", avg_loss_mlm, epoch)
+        writer.add_scalar("epoch/time_sec", epoch_time, epoch)
+
+        logger.info(
+            f"Epoch {epoch+1}/{args.epochs} | Time: {epoch_time:.2f}s | "
+            f"Total Loss: {avg_loss:.4f} | L1 (MTR): {avg_loss_mtr:.4f} | "
+            f"CE (MLM): {avg_loss_mlm:.4f} | Grad Norm: {grad_norm:.4f}"
+        )
+
+    writer.close()
+
+    logger.info("Saving checkpoint..")
+    torch.save({
+        "backbone": model_base.state_dict(),
+        "mlm_head": mlm_head.state_dict(),
+        "mtr_head": rdkit_head.state_dict(),
+    }, args.output_file)
+    logger.info("Training is finished")