Release complete SpliNet Embedder

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+splinet_embedder.ts filter=lfs diff=lfs merge=lfs -text

MANIFEST.json

@@ -0,0 +1,62 @@
+{
+  "README.md": {
+    "size_bytes": 1886,
+    "sha256": "467819dc02f633b6f433c85c72c185731ede2d646b2d016fb1e401b1df803643"
+  },
+  "api.py": {
+    "size_bytes": 3460,
+    "sha256": "c91422223e045a612ff8b84321430dbd5ee956545df71479ba96822044fd50c2"
+  },
+  "evaluation/FINAL_EVALUATION_SUMMARY.json": {
+    "size_bytes": 756,
+    "sha256": "4d3547d01b74bf724d7983f3beaa207fbee977cb41d373d33ae07133cf4fe74d"
+  },
+  "evaluation/STSBenchmark.json": {
+    "size_bytes": 85,
+    "sha256": "eee9ba254e1f81a6c0ecce23663eb9d675c91976e1557fa7a8006367b33751e1"
+  },
+  "evaluation/SciFact.json": {
+    "size_bytes": 220,
+    "sha256": "ba66fde9c4a72ebfd50d5306bb334708c788f562f11790027e9e2e65cb6be2be"
+  },
+  "evaluation/TwentyNewsgroupsClustering.json": {
+    "size_bytes": 2581,
+    "sha256": "bd71ae025dbf831dcb03ea42429d527822968e134735187c9910f9ab80e2c91f"
+  },
+  "evaluation/inference_vram_autotune.json": {
+    "size_bytes": 3081,
+    "sha256": "bb6dd07068b2fff8370c406140802d091328577d845baf7adac2e995e1f67dd6"
+  },
+  "model_info.json": {
+    "size_bytes": 1336,
+    "sha256": "9a7d6ec81e4c52246801790973abd4b626c8df1787e65019551768002edaf73f"
+  },
+  "requirements.txt": {
+    "size_bytes": 115,
+    "sha256": "bdaa62a067a12591d656fc74e410ee9552f53c8cf9b4941a0b92200e3dc83bd2"
+  },
+  "spiece.model": {
+    "size_bytes": 806898,
+    "sha256": "119ec6b2af9cbbc56f297bd606b69f79f6e3130a34ee56122ee813a58d15bb9d"
+  },
+  "splinet_embedder.ts": {
+    "size_bytes": 329392397,
+    "sha256": "7800d07a6bf69a0fa06ae72cfc553aea4ef263634ec5eb462008927110ea7aea"
+  },
+  "splinet_embedder_full_qqp.pt": {
+    "size_bytes": 329128162,
+    "sha256": "ae6a2f15ac997d3ca165de09fee89789c181dc2f62ccc6ef2820c6db5f475b05"
+  },
+  "training/TRAINING_COMPLETE.json": {
+    "size_bytes": 381,
+    "sha256": "6186b759a19b1ad12082e6c61b19884728eb2e5909366ff4a76761b39287cf9f"
+  },
+  "training/embedder_config.json": {
+    "size_bytes": 218,
+    "sha256": "3eaa07d108097256de4e5f929c32d4b4444479a2ddeb55dbc907ca021049264c"
+  },
+  "training/vram_autotune.json": {
+    "size_bytes": 2696,
+    "sha256": "1df1f89009031834701cfbd8ed7ef9735f68cbc589456a357db9073245af18d6"
+  }
+}

README.md

@@ -0,0 +1,75 @@
+---
+license: agpl-3.0
+pipeline_tag: feature-extraction
+language:
+- en
+tags:
+- sentence-embeddings
+- sentence-similarity
+- splinet
+- torchscript
+---
+
+# SpliNet Embedder
+
+Sentence-embedding version of SpliNet.
+
+The complete encoder was fine-tuned end-to-end on **134,378 positive QQP duplicate pairs** using symmetric contrastive learning.
+
+## Model
+
+| Property | Value |
+|---|---:|
+| Encoder parameters | **82,270,464** |
+| Embedding dimension | **768** |
+| Maximum input length | **512 tokens** |
+| Pooling | Masked mean |
+| Output | L2-normalized |
+| Fine-tuning pairs | **134,378** |
+| Epochs | **1** |
+| Optimization steps | **8,958** |
+
+## Full evaluation
+
+| Benchmark | Metric | Score |
+|---|---|---:|
+| STSBenchmark | Pearson | **0.5257** |
+| STSBenchmark | Spearman | **0.5174** |
+| SciFact | nDCG@10 | **0.1928** |
+| SciFact | MAP@10 | **0.1686** |
+| SciFact | MRR@10 | **0.1766** |
+| SciFact | Recall@10 | **0.2614** |
+| SciFact | Recall@100 | **0.4413** |
+| Twenty Newsgroups | Mean V-measure | **0.1532** |
+| Twenty Newsgroups | Mean homogeneity | **0.1510** |
+| Twenty Newsgroups | Mean completeness | **0.1555** |
+
+Twenty Newsgroups evaluation used **59,545 texts** across all ten subsets.
+
+## Files
+
+- `splinet_embedder.ts` — complete deployable TorchScript encoder
+- `splinet_embedder_full_qqp.pt` — original trained checkpoint
+- `spiece.model` — tokenizer
+- `api.py` — Python/FastAPI interface
+- `model_info.json` — metadata
+- `evaluation/` — benchmark results
+- `training/` — final training metadata
+
+## Usage
+
+Install `requirements.txt`, then:
+
+    from api import encode, similarity
+    z = encode(['first sentence', 'second sentence'])
+    print(z.shape)
+
+To serve:
+
+    python api.py
+
+Endpoints: `/health`, `/model-info`, `/embed`, `/similarity`.
+
+Base model: https://huggingface.co/Angshul/SpliNet
+
+Source: https://github.com/AngshulMajumdar/SpliNet

api.py

@@ -0,0 +1,106 @@
+from pathlib import Path
+import os
+import json
+import unicodedata
+import numpy as np
+import sentencepiece as spm
+import torch
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+
+ROOT = Path(__file__).resolve().parent
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+MODEL = torch.jit.load(str(ROOT / "splinet_embedder.ts"), map_location=DEVICE).eval()
+SP = spm.SentencePieceProcessor(model_file=str(ROOT / "spiece.model"))
+
+PAD_ID = 3
+CLS_ID = 4
+SEP_ID = 5
+MAX_LENGTH = 512
+DEFAULT_BATCH = int(os.environ.get("SPLINET_BATCH_SIZE", "256" if DEVICE.type == "cuda" else "8"))
+
+def normalize_text(text):
+    text = unicodedata.normalize("NFKC", text or "").lower()
+    return " ".join(text.split())
+
+def tokenize(texts):
+    ids = torch.full((len(texts), MAX_LENGTH), PAD_ID, dtype=torch.long)
+    mask = torch.zeros((len(texts), MAX_LENGTH), dtype=torch.long)
+    for i, text in enumerate(texts):
+        pieces = list(SP.encode(normalize_text(text), out_type=int))[:MAX_LENGTH - 2]
+        seq = [CLS_ID] + pieces + [SEP_ID]
+        n = len(seq)
+        ids[i, :n] = torch.tensor(seq, dtype=torch.long)
+        mask[i, :n] = 1
+    return ids, mask
+
+def encode(texts, batch_size=None):
+    if isinstance(texts, str):
+        texts = [texts]
+    if not texts:
+        return np.empty((0, 768), dtype=np.float32)
+    bs = max(1, int(batch_size or DEFAULT_BATCH))
+    output = []
+    start = 0
+    while start < len(texts):
+        current = min(bs, len(texts) - start)
+        try:
+            ids, mask = tokenize(texts[start:start + current])
+            ids = ids.to(DEVICE)
+            mask = mask.to(DEVICE)
+            with torch.inference_mode():
+                if DEVICE.type == "cuda":
+                    with torch.autocast(device_type="cuda", dtype=torch.float16):
+                        z = MODEL(ids, mask)
+                else:
+                    z = MODEL(ids, mask)
+            output.append(z.float().cpu().numpy())
+            start += current
+        except torch.OutOfMemoryError:
+            if DEVICE.type != "cuda" or current == 1:
+                raise
+            torch.cuda.empty_cache()
+            bs = max(1, current // 2)
+    return np.concatenate(output, axis=0).astype(np.float32)
+
+def similarity(text_a, text_b):
+    z = encode([text_a, text_b])
+    return float(np.dot(z[0], z[1]))
+
+class EmbedRequest(BaseModel):
+    texts: list[str]
+    batch_size: int | None = None
+
+class SimilarityRequest(BaseModel):
+    text_a: str
+    text_b: str
+
+app = FastAPI(title="SpliNet Embedder API", version="1.0.0")
+
+@app.get("/health")
+def health():
+    return {"status": "ok", "device": str(DEVICE)}
+
+@app.get("/model-info")
+def model_info():
+    with open(ROOT / "model_info.json", "r", encoding="utf-8") as f:
+        return json.load(f)
+
+@app.post("/embed")
+def embed_endpoint(req: EmbedRequest):
+    try:
+        z = encode(req.texts, req.batch_size)
+        return {"count": len(req.texts), "dimension": 768, "embeddings": z.tolist()}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/similarity")
+def similarity_endpoint(req: SimilarityRequest):
+    try:
+        return {"cosine_similarity": similarity(req.text_a, req.text_b)}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "8000")))

evaluation/FINAL_EVALUATION_SUMMARY.json

@@ -0,0 +1,28 @@
+{
+  "model": "SpliNet Sentence Embedder",
+  "base_model": "Angshul/SpliNet",
+  "embedding_dimension": 768,
+  "sequence_length": 512,
+  "inference_batch": 527,
+  "measured_peak_reserved_gb": 14.041015625,
+  "STSBenchmark": {
+    "pairs": 1379,
+    "pearson": 0.525678022791122,
+    "spearman": 0.5173543800911169
+  },
+  "SciFact": {
+    "nDCG@10": 0.19278095226166866,
+    "MAP@10": 0.16858928571428572,
+    "MRR@10": 0.17657539682539683,
+    "Recall@10": 0.2613888888888889,
+    "Recall@100": 0.44133333333333336,
+    "corpus_documents": 5183,
+    "queries": 300
+  },
+  "TwentyNewsgroups": {
+    "total_texts": 59545,
+    "mean_homogeneity": 0.15101690743103127,
+    "mean_completeness": 0.15554516291013415,
+    "mean_v_measure": 0.15324401136971566
+  }
+}

evaluation/STSBenchmark.json

@@ -0,0 +1,5 @@
+{
+  "pairs": 1379,
+  "pearson": 0.525678022791122,
+  "spearman": 0.5173543800911169
+}

evaluation/SciFact.json

@@ -0,0 +1,9 @@
+{
+  "nDCG@10": 0.19278095226166866,
+  "MAP@10": 0.16858928571428572,
+  "MRR@10": 0.17657539682539683,
+  "Recall@10": 0.2613888888888889,
+  "Recall@100": 0.44133333333333336,
+  "corpus_documents": 5183,
+  "queries": 300
+}

evaluation/TwentyNewsgroupsClustering.json

@@ -0,0 +1,99 @@
+{
+  "num_subsets": 10,
+  "total_texts": 59545,
+  "mean_homogeneity": 0.15101690743103127,
+  "mean_completeness": 0.15554516291013415,
+  "mean_v_measure": 0.15324401136971566,
+  "subsets": [
+    {
+      "subset": 0,
+      "num_texts": 1000,
+      "num_classes": 20,
+      "homogeneity": 0.18848367785332787,
+      "completeness": 0.19373099116017206,
+      "v_measure": 0.19107131509256678,
+      "inertia": 770.6087646484375
+    },
+    {
+      "subset": 1,
+      "num_texts": 2101,
+      "num_classes": 20,
+      "homogeneity": 0.16526631539460462,
+      "completeness": 0.1681524678928551,
+      "v_measure": 0.16669690003158819,
+      "inertia": 1623.0032958984375
+    },
+    {
+      "subset": 2,
+      "num_texts": 3202,
+      "num_classes": 20,
+      "homogeneity": 0.1594426468684236,
+      "completeness": 0.16370453150846842,
+      "v_measure": 0.1615454848726722,
+      "inertia": 2477.881591796875
+    },
+    {
+      "subset": 3,
+      "num_texts": 4303,
+      "num_classes": 20,
+      "homogeneity": 0.14095093771777503,
+      "completeness": 0.1450873980629758,
+      "v_measure": 0.1429892587802194,
+      "inertia": 3356.320068359375
+    },
+    {
+      "subset": 4,
+      "num_texts": 5404,
+      "num_classes": 20,
+      "homogeneity": 0.15037425141416783,
+      "completeness": 0.15393466315321286,
+      "v_measure": 0.15213362889000107,
+      "inertia": 4200.5
+    },
+    {
+      "subset": 5,
+      "num_texts": 6505,
+      "num_classes": 20,
+      "homogeneity": 0.13539065429544844,
+      "completeness": 0.1427104755759666,
+      "v_measure": 0.1389542334616076,
+      "inertia": 5078.7021484375
+    },
+    {
+      "subset": 6,
+      "num_texts": 7606,
+      "num_classes": 20,
+      "homogeneity": 0.13514684478067546,
+      "completeness": 0.1389571821370939,
+      "v_measure": 0.13702552958899597,
+      "inertia": 5929.64111328125
+    },
+    {
+      "subset": 7,
+      "num_texts": 8707,
+      "num_classes": 20,
+      "homogeneity": 0.15154795648155694,
+      "completeness": 0.15817573965897352,
+      "v_measure": 0.15479093404206168,
+      "inertia": 6761.9462890625
+    },
+    {
+      "subset": 8,
+      "num_texts": 9808,
+      "num_classes": 20,
+      "homogeneity": 0.14624533873516948,
+      "completeness": 0.1495133374423665,
+      "v_measure": 0.1478612831399018,
+      "inertia": 7655.23095703125
+    },
+    {
+      "subset": 9,
+      "num_texts": 10909,
+      "num_classes": 20,
+      "homogeneity": 0.13732045076916352,
+      "completeness": 0.1414848425092566,
+      "v_measure": 0.139371545797542,
+      "inertia": 8504.7734375
+    }
+  ]
+}

evaluation/inference_vram_autotune.json

@@ -0,0 +1,127 @@
+{
+  "gpu": "Tesla T4",
+  "total_vram_gb": 14.56317138671875,
+  "target_vram_gb": 13.5,
+  "hard_ceiling_fraction": 0.965,
+  "hard_ceiling_gb": 14.053460388183593,
+  "selected_inference_batch": 527,
+  "selected_peak_reserved_gb": 14.041015625,
+  "selected_peak_allocated_gb": 12.066949844360352,
+  "measurements": [
+    {
+      "batch_size": 64,
+      "success": true,
+      "peak_allocated_gb": 2.224550247192383,
+      "peak_reserved_gb": 2.7890625,
+      "above_target": false,
+      "below_ceiling": true
+    },
+    {
+      "batch_size": 128,
+      "success": true,
+      "peak_allocated_gb": 3.584657669067383,
+      "peak_reserved_gb": 4.666015625,
+      "above_target": false,
+      "below_ceiling": true
+    },
+    {
+      "batch_size": 256,
+      "success": true,
+      "peak_allocated_gb": 6.304872512817383,
+      "peak_reserved_gb": 8.41796875,
+      "above_target": false,
+      "below_ceiling": true
+    },
+    {
+      "batch_size": 512,
+      "success": true,
+      "peak_allocated_gb": 11.745302200317383,
+      "peak_reserved_gb": 13.666015625,
+      "above_target": true,
+      "below_ceiling": true
+    },
+    {
+      "batch_size": 520,
+      "success": true,
+      "peak_allocated_gb": 11.916261672973633,
+      "peak_reserved_gb": 13.86328125,
+      "above_target": true,
+      "below_ceiling": true
+    },
+    {
+      "batch_size": 524,
+      "success": true,
+      "peak_allocated_gb": 12.001253128051758,
+      "peak_reserved_gb": 13.962890625,
+      "above_target": true,
+      "below_ceiling": true
+    },
+    {
+      "batch_size": 526,
+      "success": true,
+      "peak_allocated_gb": 12.04667854309082,
+      "peak_reserved_gb": 14.013671875,
+      "above_target": true,
+      "below_ceiling": true
+    },
+    {
+      "batch_size": 527,
+      "success": true,
+      "peak_allocated_gb": 12.066949844360352,
+      "peak_reserved_gb": 14.041015625,
+      "above_target": true,
+      "below_ceiling": true
+    },
+    {
+      "batch_size": 527,
+      "success": true,
+      "peak_allocated_gb": 12.066949844360352,
+      "peak_reserved_gb": 14.041015625,
+      "above_target": true,
+      "below_ceiling": true
+    },
+    {
+      "batch_size": 528,
+      "success": true,
+      "peak_allocated_gb": 12.086244583129883,
+      "peak_reserved_gb": 14.0625,
+      "above_target": true,
+      "below_ceiling": false
+    },
+    {
+      "batch_size": 544,
+      "success": false,
+      "above_target": false,
+      "below_ceiling": false,
+      "error": "CUDA OOM"
+    },
+    {
+      "batch_size": 576,
+      "success": false,
+      "above_target": false,
+      "below_ceiling": false,
+      "error": "CUDA OOM"
+    },
+    {
+      "batch_size": 640,
+      "success": false,
+      "above_target": false,
+      "below_ceiling": false,
+      "error": "CUDA OOM"
+    },
+    {
+      "batch_size": 768,
+      "success": false,
+      "above_target": false,
+      "below_ceiling": false,
+      "error": "CUDA OOM"
+    },
+    {
+      "batch_size": 1024,
+      "success": false,
+      "above_target": false,
+      "below_ceiling": false,
+      "error": "CUDA OOM"
+    }
+  ]
+}

model_info.json

@@ -0,0 +1,48 @@
+{
+  "name": "SpliNet Embedder",
+  "base_model": "Angshul/SpliNet",
+  "encoder_parameters": 82270464,
+  "embedding_dimension": 768,
+  "maximum_sequence_length": 512,
+  "pooling": "masked mean",
+  "output_normalization": "L2",
+  "fine_tuning": {
+    "dataset": "QQP positive duplicate pairs",
+    "pairs": 134378,
+    "epochs": 1,
+    "steps": 8958
+  },
+  "files": {
+    "deployment_model": "splinet_embedder.ts",
+    "original_checkpoint": "splinet_embedder_full_qqp.pt",
+    "tokenizer": "spiece.model"
+  },
+  "evaluation": {
+    "model": "SpliNet Sentence Embedder",
+    "base_model": "Angshul/SpliNet",
+    "embedding_dimension": 768,
+    "sequence_length": 512,
+    "inference_batch": 527,
+    "measured_peak_reserved_gb": 14.041015625,
+    "STSBenchmark": {
+      "pairs": 1379,
+      "pearson": 0.525678022791122,
+      "spearman": 0.5173543800911169
+    },
+    "SciFact": {
+      "nDCG@10": 0.19278095226166866,
+      "MAP@10": 0.16858928571428572,
+      "MRR@10": 0.17657539682539683,
+      "Recall@10": 0.2613888888888889,
+      "Recall@100": 0.44133333333333336,
+      "corpus_documents": 5183,
+      "queries": 300
+    },
+    "TwentyNewsgroups": {
+      "total_texts": 59545,
+      "mean_homogeneity": 0.15101690743103127,
+      "mean_completeness": 0.15554516291013415,
+      "mean_v_measure": 0.15324401136971566
+    }
+  }
+}

requirements.txt

@@ -0,0 +1,7 @@
+torch>=2.6
+numpy>=1.26
+sentencepiece>=0.2
+fastapi>=0.115
+uvicorn[standard]>=0.30
+pydantic>=2
+huggingface_hub>=0.34

spiece.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:119ec6b2af9cbbc56f297bd606b69f79f6e3130a34ee56122ee813a58d15bb9d
+size 806898

splinet_embedder.ts

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7800d07a6bf69a0fa06ae72cfc553aea4ef263634ec5eb462008927110ea7aea
+size 329392397

splinet_embedder_full_qqp.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ae6a2f15ac997d3ca165de09fee89789c181dc2f62ccc6ef2820c6db5f475b05
+size 329128162

training/TRAINING_COMPLETE.json

@@ -0,0 +1,9 @@
+{
+  "training_complete": true,
+  "global_step": 8958,
+  "epochs": 1,
+  "training_pairs": 134378,
+  "final_model": "/content/drive/MyDrive/SpliNet_Embedder/final_embedder_full_qqp/splinet_embedder_full_qqp.pt",
+  "log_file": "/content/drive/MyDrive/SpliNet_Embedder/training_full_qqp.jsonl",
+  "checkpoint_directory": "/content/drive/MyDrive/SpliNet_Embedder/checkpoints_full_qqp"
+}

training/embedder_config.json

@@ -0,0 +1,9 @@
+{
+  "base_model": "Angshul/SpliNet",
+  "base_parameters": 82894592,
+  "embedding_dimension": 768,
+  "pooling": "masked_mean",
+  "normalization": "L2",
+  "required_sequence_length": 512,
+  "training_precision": "fp16"
+}

training/vram_autotune.json

@@ -0,0 +1,95 @@
+{
+  "gpu": "Tesla T4",
+  "total_vram_gb": 14.56317138671875,
+  "safety_fraction": 0.92,
+  "safe_vram_ceiling_gb": 13.39811767578125,
+  "sequence_length": 512,
+  "precision": "fp16",
+  "optimizer": "AdamW",
+  "full_encoder_finetuning": true,
+  "temperature": 0.05,
+  "optimal_micro_batch_pairs": 15,
+  "sentences_per_micro_batch": 30,
+  "measurements": [
+    {
+      "pair_batch_size": 1,
+      "sentence_batch_size": 2,
+      "success": true,
+      "within_safety_margin": true,
+      "loss": 0.0,
+      "peak_allocated_gb": 1.8602900505065918,
+      "peak_reserved_gb": 2.072265625,
+      "vram_fraction": 0.14229494180710217
+    },
+    {
+      "pair_batch_size": 2,
+      "sentence_batch_size": 4,
+      "success": true,
+      "within_safety_margin": true,
+      "loss": 0.342132568359375,
+      "peak_allocated_gb": 2.849968910217285,
+      "peak_reserved_gb": 3.099609375,
+      "vram_fraction": 0.21283889976236678
+    },
+    {
+      "pair_batch_size": 4,
+      "sentence_batch_size": 8,
+      "success": true,
+      "within_safety_margin": true,
+      "loss": 0.7950439453125,
+      "peak_allocated_gb": 4.306286811828613,
+      "peak_reserved_gb": 4.677734375,
+      "vram_fraction": 0.32120300247691774
+    },
+    {
+      "pair_batch_size": 8,
+      "sentence_batch_size": 16,
+      "success": true,
+      "within_safety_margin": true,
+      "loss": 1.4732666015625,
+      "peak_allocated_gb": 7.2357072830200195,
+      "peak_reserved_gb": 7.658203125,
+      "vram_fraction": 0.5258609489402899
+    },
+    {
+      "pair_batch_size": 12,
+      "sentence_batch_size": 24,
+      "success": true,
+      "within_safety_margin": true,
+      "loss": 1.0550334453582764,
+      "peak_allocated_gb": 10.169827461242676,
+      "peak_reserved_gb": 10.623046875,
+      "vram_fraction": 0.7294459834955973
+    },
+    {
+      "pair_batch_size": 14,
+      "sentence_batch_size": 28,
+      "success": true,
+      "within_safety_margin": true,
+      "loss": 0.9702497720718384,
+      "peak_allocated_gb": 11.684922218322754,
+      "peak_reserved_gb": 12.13671875,
+      "vram_fraction": 0.8333843245893807
+    },
+    {
+      "pair_batch_size": 15,
+      "sentence_batch_size": 30,
+      "success": true,
+      "within_safety_margin": true,
+      "loss": 0.7030029296875,
+      "peak_allocated_gb": 12.402628898620605,
+      "peak_reserved_gb": 12.9140625,
+      "vram_fraction": 0.8867616920156075
+    },
+    {
+      "pair_batch_size": 16,
+      "sentence_batch_size": 32,
+      "success": true,
+      "within_safety_margin": false,
+      "loss": 1.680816650390625,
+      "peak_allocated_gb": 13.109440803527832,
+      "peak_reserved_gb": 13.408203125,
+      "vram_fraction": 0.9206925311081587
+    }
+  ]
+}