Upload 4 files

chat.py

@@ -0,0 +1,422 @@
+import torch
+import torch.nn.functional as F
+
+from model import GenoLiteHybrid
+
+# =========================================================
+# CONFIG
+# =========================================================
+
+DEVICE = torch.device(
+    "cuda" if torch.cuda.is_available() else "cpu"
+)
+
+CHUNK_SIZE = 64
+TOKEN_MAP = {
+    "U": 0,
+    "D": 1,
+    "-": 2,
+    "+": 3,
+    "J": 4,
+    "R": 5,
+    "L": 6,
+    "T": 7,
+    "C": 8,
+    "H": 9,
+    "F": 10
+}
+
+ID2LABEL = {
+    0: "0",
+    1: "1",
+    2: "2",
+    3: "3",
+    4: "4",
+    5: "5",
+    6: "6",
+    7: "7",
+    8: "8",
+    9: "9"
+}
+
+# =========================================================
+# LOAD MODEL
+# =========================================================
+
+model = GenoLiteHybrid().to(DEVICE)
+
+checkpoint = torch.load(
+    "model.pt",
+    map_location=DEVICE
+)
+
+# ---------------------------------------------------------
+# RAW OR FULL CHECKPOINT
+# ---------------------------------------------------------
+
+if isinstance(checkpoint, dict) and \
+   "model_state_dict" in checkpoint:
+
+    model.load_state_dict(
+        checkpoint["model_state_dict"]
+    )
+
+    print("\nLoaded full checkpoint.")
+
+else:
+
+    model.load_state_dict(checkpoint)
+
+    print("\nLoaded raw state_dict.")
+
+model.eval()
+
+print("\n===================================")
+print("          MODEL LOADED")
+print("===================================\n")
+
+# =========================================================
+# ENCODE
+# =========================================================
+
+def encode(seq):
+
+    return torch.tensor(
+
+        [TOKEN_MAP[c] for c in seq],
+
+        dtype=torch.long
+    )
+
+# =========================================================
+# CHUNKING
+# =========================================================
+
+def split_chunks(sequence):
+
+    chunks = []
+
+    for i in range(
+        0,
+        len(sequence),
+        CHUNK_SIZE
+    ):
+
+        chunk = sequence[
+            i:i + CHUNK_SIZE
+        ]
+
+        chunks.append(chunk)
+
+    return chunks
+
+# =========================================================
+# SINGLE CHUNK INFERENCE
+# =========================================================
+
+def analyze_chunk(sequence):
+
+    x = encode(sequence)
+
+    x = x.unsqueeze(0).to(DEVICE)
+
+    with torch.no_grad():
+
+        # ---------------------------------------------
+        # EMBEDDING
+        # ---------------------------------------------
+
+        emb = model.embedding(x)
+
+        # ---------------------------------------------
+        # EXPERTS
+        # ---------------------------------------------
+
+        cnn_out = model.cnn(emb)
+
+        gru_out = model.gru(emb)
+
+        tf_out = model.transformer(emb)
+
+        mamba_out = model.mamba(emb)
+
+        # ---------------------------------------------
+        # EXPERT ACTIVITY
+        # ---------------------------------------------
+
+        cnn_score = cnn_out.abs().mean().item()
+
+        gru_score = gru_out.abs().mean().item()
+
+        tf_score = tf_out.abs().mean().item()
+
+        mamba_score = mamba_out.abs().mean().item()
+
+        total = (
+            cnn_score +
+            gru_score +
+            tf_score +
+            mamba_score
+        )
+
+        cnn_w = cnn_score / total
+        gru_w = gru_score / total
+        tf_w = tf_score / total
+        mamba_w = mamba_score / total
+
+        # ---------------------------------------------
+        # FINAL PRED
+        # ---------------------------------------------
+
+        fused = torch.cat(
+            [
+                cnn_out,
+                gru_out,
+                tf_out,
+                mamba_out
+            ],
+            dim=-1
+        )
+
+        fused = model.fusion(fused)
+
+        pooled = fused.mean(dim=1)
+
+        logits = model.classifier(pooled)
+
+        probs = F.softmax(
+            logits,
+            dim=-1
+        )
+
+        pred = probs.argmax(dim=-1).item()
+
+    return {
+
+        "prediction": ID2LABEL[pred],
+
+        "probs": probs[0].cpu(),
+
+        "cnn": cnn_w,
+        "gru": gru_w,
+        "tf": tf_w,
+        "mamba": mamba_w
+    }
+
+# =========================================================
+# FULL ANALYSIS
+# =========================================================
+
+def analyze_sequence(sequence):
+
+    sequence = sequence.strip().upper()
+
+    # -----------------------------------------------------
+    # VALIDATION
+    # -----------------------------------------------------
+
+    valid = all(
+        c in TOKEN_MAP
+        for c in sequence
+    )
+
+    if not valid:
+
+        print("\nOnly A/T/G/C allowed.\n")
+        return
+
+    # -----------------------------------------------------
+    # LENGTH CHECK
+    # -----------------------------------------------------
+
+    length = len(sequence)
+
+    if length < CHUNK_SIZE:
+
+        missing = CHUNK_SIZE - length
+
+        print("\n===================================")
+        print("           LENGTH ERROR")
+        print("===================================\n")
+
+        print("Input too short.\n")
+
+        print(
+            f"Current Length : {length}"
+        )
+
+        print(
+            f"Missing Chars  : {missing}"
+        )
+
+        print(
+            f"Required Length: {CHUNK_SIZE}"
+        )
+
+        print("\n===================================\n")
+
+        return
+
+    # -----------------------------------------------------
+    # MULTIPLE CHECK
+    # -----------------------------------------------------
+
+    if length % CHUNK_SIZE != 0:
+
+        next_valid = (
+            (
+                length // CHUNK_SIZE
+            ) + 1
+        ) * CHUNK_SIZE
+
+        missing = next_valid - length
+
+        print("\n===================================")
+        print("           LENGTH ERROR")
+        print("===================================\n")
+
+        print(
+            f"Sequence length must be "
+            f"a multiple of {CHUNK_SIZE}.\n"
+        )
+
+        print(
+            f"Current Length : {length}"
+        )
+
+        print(
+            f"Next Valid Size: {next_valid}"
+        )
+
+        print(
+            f"Missing Chars  : {missing}"
+        )
+
+        print("\n===================================\n")
+
+        return
+
+    # -----------------------------------------------------
+    # CHUNKING
+    # -----------------------------------------------------
+
+    chunks = split_chunks(sequence)
+
+    print("\n===================================")
+    print("         ANALYZING INPUT")
+    print("===================================\n")
+
+    print(f"Total Length : {len(sequence)}")
+
+    print(f"Chunks       : {len(chunks)}")
+
+    # -----------------------------------------------------
+    # AGGREGATION
+    # -----------------------------------------------------
+
+    total_probs = torch.zeros(10)
+
+    total_cnn = 0
+    total_gru = 0
+    total_tf = 0
+    total_mamba = 0
+
+    # -----------------------------------------------------
+    # PROCESS CHUNKS
+    # -----------------------------------------------------
+
+    for idx, chunk in enumerate(chunks):
+
+        result = analyze_chunk(chunk)
+
+        total_probs += result["probs"]
+
+        total_cnn += result["cnn"]
+        total_gru += result["gru"]
+        total_tf += result["tf"]
+        total_mamba += result["mamba"]
+
+        print("\n-----------------------------------")
+        print(f"Chunk {idx+1}")
+        print("-----------------------------------\n")
+
+        print(chunk)
+
+        print("\nPrediction:")
+        print(result["prediction"])
+
+        print("\nProbabilities:\n")
+
+        for i in range(3):
+
+            print(
+                f"{ID2LABEL[i]}: "
+                f"{result['probs'][i].item():.4f}"
+            )
+
+    # -----------------------------------------------------
+    # AVERAGES
+    # -----------------------------------------------------
+
+    total_probs /= len(chunks)
+
+    total_cnn /= len(chunks)
+    total_gru /= len(chunks)
+    total_tf /= len(chunks)
+    total_mamba /= len(chunks)
+
+    # -----------------------------------------------------
+    # FINAL DECISION
+    # -----------------------------------------------------
+
+    final_pred = total_probs.argmax().item()
+
+    print("\n===================================")
+    print("          FINAL RESULT")
+    print("===================================\n")
+
+    print(
+        f"FINAL DECISION: "
+        f"{ID2LABEL[final_pred]}"
+    )
+
+    print("\n-----------------------------------")
+    print("Average Probabilities")
+    print("-----------------------------------\n")
+
+    for i in range(3):
+
+        print(
+            f"{ID2LABEL[i]}: "
+            f"{total_probs[i].item():.4f}"
+        )
+
+    print("\n-----------------------------------")
+    print("Average Expert Activity")
+    print("-----------------------------------\n")
+
+    print(f"CNN         : {total_cnn:.4f}")
+    print(f"GRU         : {total_gru:.4f}")
+    print(f"Transformer : {total_tf:.4f}")
+    print(f"Mamba       : {total_mamba:.4f}")
+
+    print("\n===================================\n")
+
+# =========================================================
+# CHAT LOOP
+# =========================================================
+
+print("Type DNA sequence.")
+print("Length must be 64 or multiples of 64.")
+print("Type EXIT to quit.\n")
+
+while True:
+
+    seq = input("logs > ")
+
+    if seq.strip().upper() == "EXIT":
+
+        print("\nBye.\n")
+        break
+
+    analyze_sequence(seq)

config.json

@@ -0,0 +1,113 @@
+{
+  "model_type": "GenoLiteHybrid",
+
+  "vocab_size": 4,
+  "sequence_length": 64,
+  "num_classes": 3,
+
+  "d_model": 512,
+
+  "cnn": {
+    "enabled": true,
+    "blocks": 7,
+    "channels": 960,
+    "kernel_size": 3,
+    "residual": true,
+    "layernorm": true,
+    "activation": "gelu"
+  },
+
+  "gru": {
+    "enabled": true,
+    "hidden_size": 960,
+    "layers": 4,
+    "bidirectional": false,
+    "batch_first": true,
+    "projection_to_d_model": true,
+    "layernorm": true
+  },
+
+  "transformer": {
+    "enabled": true,
+    "layers": 6,
+    "heads": 8,
+    "ffn_dim": 2048,
+    "dropout": 0.1,
+    "activation": "gelu",
+    "batch_first": true,
+    "layernorm": true
+  },
+
+  "mamba": {
+    "enabled": true,
+    "layers": 10,
+    "state_dim": 1408,
+    "gated": true,
+    "residual": true,
+    "layernorm": true
+  },
+
+  "fusion": {
+    "input_dim": 2048,
+    "output_dim": 512,
+    "activation": "gelu",
+    "dropout": 0.1,
+    "layernorm": true
+  },
+
+  "classifier": {
+    "hidden_dim": 512,
+    "dropout": 0.1,
+    "activation": "gelu",
+    "num_classes": 3
+  },
+
+  "pooling": {
+    "type": "mean"
+  },
+
+  "training": {
+    "epochs": 3,
+    "batch_size": 3,
+    "learning_rate": 0.0001,
+    "optimizer": "AdamW",
+    "weight_decay": 0.01,
+    "gradient_clipping": 1.0,
+    "shuffle": true
+  },
+
+  "dataset": {
+    "type": "synthetic",
+    "samples_total": 9000,
+    "samples_per_class": 3000,
+
+    "classes": [
+      "OK",
+      "MHAP",
+      "PROBLEM"
+    ],
+
+    "difficulty_levels": [
+      "easy",
+      "medium",
+      "hard"
+    ],
+
+    "features": [
+      "controlled_entropy",
+      "motif_repetition",
+      "hidden_illegal_pairs",
+      "partial_shuffle",
+      "duplicate_prevention",
+      "class_overlap"
+    ]
+  },
+
+  "hardware": {
+    "device": "cpu",
+    "ram_gb": 8,
+    "cpu": "Intel i7-4700MQ"
+  },
+
+  "estimated_parameters": "88M+"
+}

model.pt

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

model.py

@@ -0,0 +1,401 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# =========================================================
+# CONFIG
+# =========================================================
+
+VOCAB_SIZE = 11
+SEQ_LEN = 64
+NUM_CLASSES = 10
+
+D_MODEL = 512
+
+CONFIG = {
+
+    # -----------------------------------------------------
+    # CNN
+    # -----------------------------------------------------
+
+    "cnn": {
+        "blocks": 7,
+        "channels": 960,
+        "kernel": 3
+    },
+
+    # -----------------------------------------------------
+    # GRU
+    # -----------------------------------------------------
+
+    "gru": {
+        "hidden": 960,
+        "layers": 4
+    },
+
+    # -----------------------------------------------------
+    # TRANSFORMER
+    # -----------------------------------------------------
+
+    "transformer": {
+        "layers": 6,
+        "heads": 8,
+        "ffn": 2048,
+        "dropout": 0.1
+    },
+
+    # -----------------------------------------------------
+    # MAMBA-LIKE
+    # -----------------------------------------------------
+
+    "mamba": {
+        "layers": 10,
+        "state_dim": 1408
+    }
+}
+
+# =========================================================
+# CNN EXPERT
+# =========================================================
+
+class CNNBlock(nn.Module):
+    def __init__(self, channels, kernel):
+        super().__init__()
+
+        self.conv1 = nn.Conv1d(
+            D_MODEL,
+            channels,
+            kernel_size=kernel,
+            padding=kernel // 2
+        )
+
+        self.conv2 = nn.Conv1d(
+            channels,
+            D_MODEL,
+            kernel_size=kernel,
+            padding=kernel // 2
+        )
+
+        self.norm = nn.LayerNorm(D_MODEL)
+
+    def forward(self, x):
+
+        # x = [B, S, D]
+
+        residual = x
+
+        x = x.transpose(1, 2)  # [B, D, S]
+
+        x = self.conv1(x)
+        x = F.gelu(x)
+
+        x = self.conv2(x)
+        x = F.gelu(x)
+
+        x = x.transpose(1, 2)  # [B, S, D]
+
+        x = x + residual
+
+        return self.norm(x)
+
+class CNNExpert(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.blocks = nn.ModuleList([
+            CNNBlock(
+                channels=config["channels"],
+                kernel=config["kernel"]
+            )
+            for _ in range(config["blocks"])
+        ])
+
+        self.norm = nn.LayerNorm(D_MODEL)
+
+    def forward(self, x):
+
+        for block in self.blocks:
+            x = block(x)
+
+        return self.norm(x)
+
+# =========================================================
+# GRU EXPERT
+# =========================================================
+
+class GRUExpert(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.gru = nn.GRU(
+            input_size=D_MODEL,
+            hidden_size=config["hidden"],
+            num_layers=config["layers"],
+            batch_first=True
+        )
+
+        self.proj = nn.Linear(
+            config["hidden"],
+            D_MODEL
+        )
+
+        self.norm = nn.LayerNorm(D_MODEL)
+
+    def forward(self, x):
+
+        x, _ = self.gru(x)
+
+        x = self.proj(x)
+
+        return self.norm(x)
+
+# =========================================================
+# TRANSFORMER EXPERT
+# =========================================================
+
+class TransformerExpert(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=D_MODEL,
+            nhead=config["heads"],
+            dim_feedforward=config["ffn"],
+            dropout=config["dropout"],
+            batch_first=True,
+            activation="gelu"
+        )
+
+        self.encoder = nn.TransformerEncoder(
+            encoder_layer,
+            num_layers=config["layers"]
+        )
+
+        self.norm = nn.LayerNorm(D_MODEL)
+
+    def forward(self, x):
+
+        x = self.encoder(x)
+
+        return self.norm(x)
+
+# =========================================================
+# MAMBA-LIKE BLOCK
+# =========================================================
+
+class MambaLikeBlock(nn.Module):
+    def __init__(self, state_dim):
+        super().__init__()
+
+        self.in_proj = nn.Linear(
+            D_MODEL,
+            state_dim
+        )
+
+        self.gate = nn.Linear(
+            D_MODEL,
+            state_dim
+        )
+
+        self.out_proj = nn.Linear(
+            state_dim,
+            D_MODEL
+        )
+
+        self.norm = nn.LayerNorm(D_MODEL)
+
+    def forward(self, x):
+
+        residual = x
+
+        h = self.in_proj(x)
+
+        g = torch.sigmoid(
+            self.gate(x)
+        )
+
+        x = h * g
+
+        x = self.out_proj(x)
+
+        x = x + residual
+
+        return self.norm(x)
+
+class MambaExpert(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.blocks = nn.ModuleList([
+            MambaLikeBlock(
+                state_dim=config["state_dim"]
+            )
+            for _ in range(config["layers"])
+        ])
+
+        self.norm = nn.LayerNorm(D_MODEL)
+
+    def forward(self, x):
+
+        for block in self.blocks:
+            x = block(x)
+
+        return self.norm(x)
+
+# =========================================================
+# HYBRID MODEL
+# =========================================================
+
+class GenoLiteHybrid(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        # -------------------------------------------------
+        # EMBEDDING
+        # -------------------------------------------------
+
+        self.embedding = nn.Embedding(
+            VOCAB_SIZE,
+            D_MODEL
+        )
+
+        # -------------------------------------------------
+        # EXPERTS
+        # -------------------------------------------------
+
+        self.cnn = CNNExpert(CONFIG["cnn"])
+
+        self.gru = GRUExpert(CONFIG["gru"])
+
+        self.transformer = TransformerExpert(
+            CONFIG["transformer"]
+        )
+
+        self.mamba = MambaExpert(CONFIG["mamba"])
+
+        # -------------------------------------------------
+        # FUSION
+        # -------------------------------------------------
+
+        self.fusion = nn.Sequential(
+
+            nn.Linear(
+                D_MODEL * 4,
+                D_MODEL
+            ),
+
+            nn.GELU(),
+
+            nn.Dropout(0.1),
+
+            nn.LayerNorm(D_MODEL)
+        )
+
+        # -------------------------------------------------
+        # CLASSIFIER
+        # -------------------------------------------------
+
+        self.classifier = nn.Sequential(
+
+            nn.Linear(
+                D_MODEL,
+                512
+            ),
+
+            nn.GELU(),
+
+            nn.Dropout(0.1),
+
+            nn.Linear(
+                512,
+                NUM_CLASSES
+            )
+        )
+
+    def forward(self, x):
+
+        # -------------------------------------------------
+        # EMBEDDING
+        # -------------------------------------------------
+
+        x = self.embedding(x)
+
+        # -------------------------------------------------
+        # EXPERTS
+        # -------------------------------------------------
+
+        cnn_out = self.cnn(x)
+
+        gru_out = self.gru(x)
+
+        tf_out = self.transformer(x)
+
+        mamba_out = self.mamba(x)
+
+        # -------------------------------------------------
+        # FUSION
+        # -------------------------------------------------
+
+        fused = torch.cat(
+            [
+                cnn_out,
+                gru_out,
+                tf_out,
+                mamba_out
+            ],
+            dim=-1
+        )
+
+        fused = self.fusion(fused)
+
+        # -------------------------------------------------
+        # GLOBAL POOLING
+        # -------------------------------------------------
+
+        pooled = fused.mean(dim=1)
+
+        # -------------------------------------------------
+        # CLASSIFIER
+        # -------------------------------------------------
+
+        logits = self.classifier(pooled)
+
+        return logits
+
+# =========================================================
+# PARAM COUNTER
+# =========================================================
+
+def count_params(model):
+    return sum(
+        p.numel()
+        for p in model.parameters()
+    )
+
+# =========================================================
+# TEST
+# =========================================================
+
+if __name__ == "__main__":
+
+    model = GenoLiteHybrid()
+
+    x = torch.randint(
+        0,
+        11,
+        (2, 64)
+    )
+
+    y = model(x)
+
+    print("\n================ TEST ================\n")
+
+    print("Input shape :", x.shape)
+
+    print("Output shape:", y.shape)
+
+    total_params = count_params(model)
+
+    print(f"\nTotal Params: {total_params / 1e6:.2f}M")
+
+    print("\n======================================\n")