Upload Healthcare-SLM v2

README.md

@@ -0,0 +1,34 @@
+---
+language: [en]
+license: mit
+tags:
+- healthcare
+- medical
+- clinical
+- slm
+- llama-style
+- rope
+- 1m-context
+- from-scratch
+pipeline_tag: text-generation
+---
+
+# Healthcare-SLM: Healthcare Small Language Model
+
+A **LLaMA-style transformer** (~33.9M params) trained from scratch on Healthcare domain data.
+Supports up to **1M token context** via RoPE.
+
+## Architecture
+| Component | Value |
+|-----------|-------|
+| Architecture | LLaMA-style (RoPE + RMSNorm + SwiGLU) |
+| Parameters | ~33.9M |
+| Layers | 8 |
+| Heads | 8 |
+| Embedding | 512 |
+| Max Context | 1,000,000 tokens |
+| Vocab | 16,000 BPE |
+| Best Loss | 0.8301405601203442 |
+
+## License
+MIT — Built from scratch.

chat.py

@@ -0,0 +1,100 @@
+#!/usr/bin/env python3
+"""Interactive chat and demo inference for Industry SLM."""
+
+import torch
+from tokenizers import Tokenizer
+from config import cfg
+from model import IndustrySLM
+
+
+def load_model(checkpoint_name="best_model.pt"):
+    device = torch.device(cfg.device)
+    ckpt_path = cfg.checkpoint_dir / checkpoint_name
+    if not ckpt_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
+
+    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+    for key, val in ckpt.get("config", {}).items():
+        if hasattr(cfg, key):
+            setattr(cfg, key, val)
+
+    model = IndustrySLM()
+    model.load_state_dict(ckpt["model_state_dict"], strict=False)
+    model = model.to(device)
+    model.eval()
+
+    tok_path = cfg.tokenizer_dir / cfg.tokenizer_filename
+    tokenizer = Tokenizer.from_file(str(tok_path))
+    print(f"Model loaded: {model.count_parameters()/1e6:.2f}M params")
+    return model, tokenizer, device
+
+
+def generate_response(model, tokenizer, device, prompt, max_tokens=None,
+                      temperature=0.8, top_k=50, top_p=0.9):
+    max_tokens = max_tokens or cfg.max_new_tokens
+    encoded = tokenizer.encode(prompt)
+    ids = encoded.ids
+    if ids and ids[-1] == 3:
+        ids = ids[:-1]
+    input_ids = torch.tensor([ids], dtype=torch.long, device=device)
+    input_len = input_ids.shape[1]
+
+    with torch.no_grad():
+        output_ids = model.generate(input_ids, max_new_tokens=max_tokens,
+                                    temperature=temperature, top_k=top_k, top_p=top_p)
+
+    new_tokens = output_ids[0][input_len:].tolist()
+    response = tokenizer.decode(new_tokens)
+    response = response.replace("<eos>", "").replace("<bos>", "").replace("<pad>", "").strip()
+    return response
+
+
+DEMO_PROMPTS = [
+    "The human cardiovascular system is",
+    "Clinical trials are essential because",
+    "Electronic health records improve patient care by",
+    "Antibiotic resistance is a growing concern because",
+    "Mental health treatment approaches include",
+]
+
+
+def demo_generation(model, tokenizer, device):
+    print(f"\n{'='*60}")
+    print(f"Demo: {cfg.domain_name}-SLM Inference")
+    print(f"{'='*60}\n")
+    for i, prompt in enumerate(DEMO_PROMPTS, 1):
+        print(f"[{i}] Prompt: {prompt}")
+        response = generate_response(model, tokenizer, device, prompt, max_tokens=256)
+        print(f"    Response: {response}\n")
+
+
+def interactive_chat():
+    print("Loading model...")
+    model, tokenizer, device = load_model()
+    print(f"\n{'='*60}")
+    print(f"{cfg.domain_name}-SLM Interactive Chat (type 'quit' to exit, 'demo' for demos)")
+    print(f"{'='*60}\n")
+    while True:
+        try:
+            user_input = input("You: ").strip()
+            if not user_input:
+                continue
+            if user_input.lower() == "quit":
+                break
+            if user_input.lower() == "demo":
+                demo_generation(model, tokenizer, device)
+                continue
+            response = generate_response(model, tokenizer, device, user_input)
+            print(f"{cfg.domain_name}-SLM: {response}\n")
+        except KeyboardInterrupt:
+            break
+    print("\nGoodbye!")
+
+
+if __name__ == "__main__":
+    import sys
+    if len(sys.argv) > 1 and sys.argv[1] == "demo":
+        model, tokenizer, device = load_model()
+        demo_generation(model, tokenizer, device)
+    else:
+        interactive_chat()

config.json

@@ -0,0 +1,18 @@
+{
+  "architectures": [
+    "IndustrySLM"
+  ],
+  "model_type": "healthcare-slm",
+  "domain": "Healthcare",
+  "vocab_size": 16000,
+  "n_layer": 8,
+  "n_head": 8,
+  "n_embd": 512,
+  "block_size": 512,
+  "dropout": 0.1,
+  "bias": false,
+  "ffn_multiplier": 2.667,
+  "max_position_embeddings": 1000000,
+  "rope_theta": 500000.0,
+  "n_parameters": 33890816
+}

config.py

@@ -0,0 +1,103 @@
+"""
+Configuration for Healthcare-SLM: A Small Language Model for Healthcare & Medical domain.
+LLaMA-style architecture with RoPE — supports up to 1M token context.
+"""
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+
+@dataclass
+class SLMConfig:
+    """All hyperparameters and paths in one place."""
+
+    # ── Project paths ──────────────────────────────────────────────
+    project_dir: Path = Path(__file__).resolve().parent
+    data_dir: Path = field(default=None)
+    tokenizer_dir: Path = field(default=None)
+    checkpoint_dir: Path = field(default=None)
+
+    # ── Domain ─────────────────────────────────────────────────────
+    domain_name: str = "Healthcare"
+    domain_slug: str = "healthcare"
+    tokenizer_filename: str = "healthcare_tokenizer.json"
+
+    # ── Tokenizer ──────────────────────────────────────────────────
+    vocab_size: int = 16_000
+    min_frequency: int = 2
+    special_tokens: list = field(
+        default_factory=lambda: [
+            "<pad>", "<unk>", "<bos>", "<eos>",
+            "<|system|>", "<|user|>", "<|assistant|>",
+        ]
+    )
+
+    # ── Model (LLaMA-style with RoPE) ─────────────────────────────
+    n_layer: int = 8
+    n_head: int = 8
+    n_embd: int = 512
+    block_size: int = 512
+    dropout: float = 0.1
+    bias: bool = False
+    ffn_multiplier: float = 2.667
+
+    # ── RoPE ───────────────────────────────────────────────────────
+    max_position_embeddings: int = 1_000_000
+    rope_theta: float = 500000.0
+
+    # ── Sliding Window ─────────────────────────────────────────────
+    sliding_window: Optional[int] = None
+
+    # ── Training ───────────────────────────────────────────────────
+    batch_size: int = 4
+    gradient_accumulation_steps: int = 4
+    learning_rate: float = 3e-4
+    weight_decay: float = 0.1
+    max_epochs: int = 20
+    dataset_stride: int = 256
+    warmup_steps: int = 20
+    grad_clip: float = 1.0
+    eval_interval: int = 50
+    eval_samples: int = 20
+    log_interval: int = 10
+    device: str = "auto"
+
+    # ── Generation ─────────────────────────────────────────────────
+    max_new_tokens: int = 1024
+    temperature: float = 0.8
+    top_k: int = 50
+    top_p: float = 0.9
+
+    # ── Hugging Face ───────────────────────────────────────────────
+    hf_repo_name: str = "healthcare-slm-1m"
+    hf_model_card_tags: list = field(
+        default_factory=lambda: [
+            "healthcare", "medical", "clinical", "slm",
+            "llama-style", "rope", "1m-context", "from-scratch",
+        ]
+    )
+
+    def __post_init__(self):
+        if self.data_dir is None:
+            self.data_dir = self.project_dir / "data"
+        if self.tokenizer_dir is None:
+            self.tokenizer_dir = self.project_dir / "tokenizer"
+        if self.checkpoint_dir is None:
+            self.checkpoint_dir = self.project_dir / "checkpoints"
+
+        self.data_dir.mkdir(parents=True, exist_ok=True)
+        self.tokenizer_dir.mkdir(parents=True, exist_ok=True)
+        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+
+        if self.device == "auto":
+            import torch
+            if torch.cuda.is_available():
+                self.device = "cuda"
+            elif torch.backends.mps.is_available():
+                self.device = "mps"
+            else:
+                self.device = "cpu"
+
+
+cfg = SLMConfig()

model.py

@@ -0,0 +1,207 @@
+#!/usr/bin/env python3
+"""
+model.py — Industry SLM Transformer with RoPE (Rotary Position Embeddings)
+===========================================================================
+Supports context lengths up to 1M tokens via:
+  * RoPE (no fixed position embedding table)
+  * RMSNorm (more efficient than LayerNorm)
+  * SwiGLU activation (better training dynamics)
+  * Flash Attention via PyTorch scaled_dot_product_attention
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple
+from config import cfg
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        norm = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+        return (x.float() * norm).type_as(x) * self.weight
+
+
+def precompute_rope_freqs(dim, max_seq_len, theta=10000.0, device=None):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, device=device).float() / dim))
+    t = torch.arange(max_seq_len, device=device).float()
+    freqs = torch.outer(t, freqs)
+    return freqs.cos(), freqs.sin()
+
+
+def apply_rope(x, cos, sin):
+    seq_len = x.shape[2]
+    head_dim = x.shape[3]
+    cos = cos[:seq_len].unsqueeze(0).unsqueeze(0)
+    sin = sin[:seq_len].unsqueeze(0).unsqueeze(0)
+    x1 = x[..., :head_dim // 2]
+    x2 = x[..., head_dim // 2:]
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        assert cfg.n_embd % cfg.n_head == 0
+        self.n_head = cfg.n_head
+        self.head_dim = cfg.n_embd // cfg.n_head
+        self.q_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False)
+        self.k_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False)
+        self.v_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False)
+        self.out_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False)
+        self.resid_drop = nn.Dropout(cfg.dropout)
+
+    def forward(self, x, rope_cos, rope_sin):
+        B, T, C = x.shape
+        q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        q = apply_rope(q, rope_cos, rope_sin)
+        k = apply_rope(k, rope_cos, rope_sin)
+        if hasattr(F, 'scaled_dot_product_attention'):
+            y = F.scaled_dot_product_attention(q, k, v,
+                dropout_p=cfg.dropout if self.training else 0.0, is_causal=True)
+        else:
+            scale = 1.0 / math.sqrt(self.head_dim)
+            att = (q @ k.transpose(-2, -1)) * scale
+            mask = torch.triu(torch.ones(T, T, device=x.device), diagonal=1).bool()
+            att = att.masked_fill(mask.unsqueeze(0).unsqueeze(0), float('-inf'))
+            att = F.softmax(att, dim=-1)
+            y = att @ v
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+        return self.resid_drop(self.out_proj(y))
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        hidden_dim = int(cfg.n_embd * getattr(cfg, 'ffn_multiplier', 2.667))
+        hidden_dim = ((hidden_dim + 63) // 64) * 64
+        self.gate_proj = nn.Linear(cfg.n_embd, hidden_dim, bias=False)
+        self.up_proj = nn.Linear(cfg.n_embd, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, cfg.n_embd, bias=False)
+        self.dropout = nn.Dropout(cfg.dropout)
+
+    def forward(self, x):
+        return self.dropout(self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn_norm = RMSNorm(cfg.n_embd)
+        self.attn = CausalSelfAttention()
+        self.ffn_norm = RMSNorm(cfg.n_embd)
+        self.ffn = SwiGLUFFN()
+
+    def forward(self, x, rope_cos, rope_sin):
+        x = x + self.attn(self.attn_norm(x), rope_cos, rope_sin)
+        x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+
+class IndustrySLM(nn.Module):
+    """Industry Small Language Model — LLaMA-style architecture with RoPE."""
+
+    def __init__(self):
+        super().__init__()
+        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd)
+        self.drop = nn.Dropout(cfg.dropout)
+        self.blocks = nn.ModuleList([TransformerBlock() for _ in range(cfg.n_layer)])
+        self.norm = RMSNorm(cfg.n_embd)
+        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
+        self.tok_emb.weight = self.lm_head.weight
+
+        head_dim = cfg.n_embd // cfg.n_head
+        max_pos = getattr(cfg, 'max_position_embeddings', 1_000_000)
+        rope_theta = getattr(cfg, 'rope_theta', 10000.0)
+        precompute_len = min(max_pos, cfg.block_size * 2)
+        cos, sin = precompute_rope_freqs(head_dim, precompute_len, theta=rope_theta)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+        self._rope_max_len = precompute_len
+        self._rope_theta = rope_theta
+        self._head_dim = head_dim
+        self.apply(self._init_weights)
+
+        n_params = sum(p.numel() for p in self.parameters())
+        print(f"{cfg.domain_name}-SLM initialized: {n_params/1e6:.2f}M parameters")
+        print(f"   Architecture: {cfg.n_layer}L / {cfg.n_head}H / {cfg.n_embd}D")
+        print(f"   Max context: {max_pos:,} tokens (via RoPE)")
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def _extend_rope(self, seq_len, device):
+        if seq_len > self._rope_max_len:
+            new_len = max(seq_len, self._rope_max_len * 2)
+            cos, sin = precompute_rope_freqs(self._head_dim, new_len,
+                                             theta=self._rope_theta, device=device)
+            self.rope_cos = cos
+            self.rope_sin = sin
+            self._rope_max_len = new_len
+
+    def forward(self, idx, targets=None):
+        B, T = idx.shape
+        device = idx.device
+        self._extend_rope(T, device)
+        x = self.drop(self.tok_emb(idx))
+        rope_cos = self.rope_cos[:T].to(device)
+        rope_sin = self.rope_sin[:T].to(device)
+        for block in self.blocks:
+            x = block(x, rope_cos, rope_sin)
+        x = self.norm(x)
+        logits = self.lm_head(x)
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        return logits, loss
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=0.8, top_k=50, top_p=0.9):
+        for _ in range(max_new_tokens):
+            idx_cond = idx if idx.size(1) <= cfg.block_size else idx[:, -cfg.block_size:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :]
+            if temperature == 0:
+                idx_next = logits.argmax(dim=-1, keepdim=True)
+            else:
+                logits = logits / temperature
+                if top_k > 0:
+                    v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                    logits[logits < v[:, [-1]]] = float('-inf')
+                if top_p < 1.0:
+                    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                    sorted_indices_to_remove = cumulative_probs > top_p
+                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                    sorted_indices_to_remove[..., 0] = 0
+                    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+                    logits[indices_to_remove] = float('-inf')
+                probs = F.softmax(logits, dim=-1)
+                idx_next = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat([idx, idx_next], dim=1)
+            if idx_next.item() == 3:  # <eos>
+                break
+        return idx
+
+    def count_parameters(self):
+        return sum(p.numel() for p in self.parameters())
+
+
+if __name__ == "__main__":
+    model = IndustrySLM()
+    x = torch.randint(0, cfg.vocab_size, (2, 32))
+    logits, loss = model(x, x)
+    print(f"Test forward: logits={logits.shape}, loss={loss.item():.4f}")

model.safetensors

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

pytorch_model.bin

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

tokenizer_config.json

@@ -0,0 +1,8 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "bos_token": "<bos>",
+  "eos_token": "<eos>",
+  "unk_token": "<unk>",
+  "pad_token": "<pad>",
+  "model_max_length": 1000000
+}