Add train.py, README.md, config.json, training_args.json, model.safetensors

README.md

@@ -0,0 +1,81 @@
+---
+license: apache-2.0
+tags:
+- pytorch
+- dino
+- multitask
+- wip
+- smoke-test
+---
+
+# Dino for Multitask
+
+## Overview
+
+Working implementation of **Dino** for **Multitask** using a **tiny** configuration. The repository focuses on transparent code and repeatable smoke tests; benchmark claims are deliberately omitted.
+
+---
+
+## Acknowledgments
+
+This repo borrows ideas freely from public reference implementations; credit to their authors.
+
+---
+
+## Repository Status
+
+- Single-file implementation; the `__main__` block runs a smoke test.
+- `config.json` and `training_args.json` hold the generated settings.
+- `model.safetensors` is freshly initialized — plumbing, not results.
+- No benchmark numbers are reported on purpose.
+
+---
+
+## Limitations
+
+- weights are random-init, so outputs carry no signal
+- robustness, fairness and domain transfer are untested
+- the API may change without notice while the design settles
+
+---
+
+## Environment
+
+| Item | Value |
+|---|---|
+| Python | >= 3.9 |
+| torch | 2.x |
+| safetensors | any recent |
+| transformers | optional, for loaders |
+
+---
+
+## Default Experiment Recipe
+
+The generated recipe uses lion with a step schedule.
+
+These values exist so the script runs end-to-end. A meaningful experiment needs its own tuning budget and at least three seeds.
+
+---
+
+## Quick Check
+
+Run `python train.py --help` for the built-in smoke test. Because this is a custom implementation, generic automatic loading APIs require an explicit adapter before use.
+
+---
+
+## Files
+
+| File | Purpose |
+|---|---|
+| `train.py` | primary artifact |
+| `README.md` | this documentation |
+| `config.json` | architecture configuration |
+| `training_args.json` | default experiment settings |
+| `model.safetensors` | initialization checkpoint |
+
+---
+
+## License
+
+Distributed under the **apache-2.0** license. External datasets keep their own terms — check them separately.

config.json

@@ -0,0 +1,19 @@
+{
+  "architectures": [
+    "CustomResearchModel"
+  ],
+  "architecture": "dino",
+  "model_type": "dino",
+  "hidden_size": 256,
+  "num_hidden_layers": 8,
+  "num_attention_heads": 4,
+  "intermediate_size": 512,
+  "hidden_act": "swish",
+  "max_position_embeddings": 128,
+  "layer_norm_eps": 1e-12,
+  "checkpoint_status": "initialization-only",
+  "notes": "Untrained checkpoint for smoke tests; no benchmark claim.",
+  "use_cache": true,
+  "torch_dtype": "float32",
+  "vocab_size": 2048
+}

model.safetensors

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

train.py

@@ -0,0 +1,126 @@
+import torch
+import torch.nn as nn
+import math
+
+
+class DinoModel(nn.Module):
+    '''
+    dino model with flash attention.
+    Scale: tiny (dim=128, layers=3, heads=2)
+    Fusion: tucker, Task: multitask
+    '''
+
+    def __init__(self, embed_dim=128, num_layers=3, num_heads=2, num_classes=10):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+
+        # image patch embedding
+        self.patch_proj = nn.Conv2d(3, embed_dim, kernel_size=16, stride=16)
+        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)
+        self.pos_embed = nn.Parameter(torch.randn(1, 197, embed_dim) * 0.02)
+        self.dropout = nn.Dropout(0.1)
+
+        # image transformer blocks
+        self.image_blocks = nn.ModuleList([
+            nn.TransformerEncoderLayer(
+                embed_dim, num_heads, embed_dim * 4, 0.1,
+                activation='gelu', batch_first=True, norm_first=True
+            )
+            for _ in range(num_layers)
+        ])
+        self.image_norm = nn.LayerNorm(embed_dim)
+
+        # text embedding
+        self.text_embed = nn.Embedding(30522, embed_dim, padding_idx=0)
+        self.text_pos = nn.Parameter(torch.randn(1, 128, embed_dim) * 0.02)
+        self.text_blocks = nn.ModuleList([
+            nn.TransformerEncoderLayer(
+                embed_dim, num_heads, embed_dim * 4, 0.1,
+                activation='gelu', batch_first=True, norm_first=True
+            )
+            for _ in range(num_layers)
+        ])
+        self.text_norm = nn.LayerNorm(embed_dim)
+
+        # fusion
+        self.fusion_blocks = nn.ModuleList([
+            nn.TransformerEncoderLayer(
+                embed_dim, num_heads, embed_dim * 4, 0.1,
+                activation='gelu', batch_first=True, norm_first=True
+            )
+            for _ in range(2)
+        ])
+        self.fusion_norm = nn.LayerNorm(embed_dim)
+
+        # task head
+        self.classifier = nn.Sequential(
+            nn.Linear(embed_dim, embed_dim),
+            nn.SiLU(),
+            nn.Dropout(0.1),
+            nn.Linear(embed_dim, num_classes),
+        )
+
+        self._initialize_weights()
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.kaiming_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Embedding):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.padding_idx is not None:
+                    m.weight[m.padding_idx].zero_()
+            elif isinstance(m, nn.LayerNorm):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+
+    def encode_image(self, images):
+        x = self.patch_proj(images)
+        x = x.flatten(2).transpose(1, 2)
+        cls = self.cls_token.expand(x.size(0), -1, -1)
+        x = torch.cat([cls, x], dim=1)
+        x = x + self.pos_embed
+        x = self.dropout(x)
+        for block in self.image_blocks:
+            x = block(x)
+        return self.image_norm(x)
+
+    def encode_text(self, input_ids, attention_mask=None):
+        x = self.text_embed(input_ids)
+        x = x + self.text_pos[:, :input_ids.size(1)]
+        x = self.dropout(x)
+        padding = (attention_mask == 0) if attention_mask is not None else None
+        for block in self.text_blocks:
+            x = block(x, src_key_padding_mask=padding)
+        return self.text_norm(x)
+
+    def forward(self, images, input_ids, attention_mask=None, labels=None):
+        image_features = self.encode_image(images)
+        text_features = self.encode_text(input_ids, attention_mask)
+
+        fused = text_features
+        for block in self.fusion_blocks:
+            fused = block(fused)
+        fused = self.fusion_norm(fused[:, 0])
+        logits = self.classifier(fused)
+
+        loss = None
+        if labels is not None:
+            loss = nn.functional.cross_entropy(logits, labels)
+
+        return {'logits': logits, 'loss': loss}
+
+
+if __name__ == '__main__':
+    model = DinoModel()
+    total = sum(p.numel() for p in model.parameters())
+    print(f'DinoModel: {total:,} params ({total/1e6:.2f}M)')
+    img = torch.randn(2, 3, 224, 224)
+    ids = torch.randint(0, 30522, (2, 128))
+    mask = torch.ones(2, 128)
+    out = model(img, ids, mask, torch.tensor([0, 1]))
+    print(f'Output: {out["logits"].shape}, Loss: {out["loss"].item():.4f}')

training_args.json

@@ -0,0 +1,9 @@
+{
+  "optimizer": "lion",
+  "scheduler": "step",
+  "learning_rate": 0.0002,
+  "batch_size": 32,
+  "epochs": 30,
+  "seed": 123,
+  "status": "default recipe; not a completed run"
+}