FSI_ECHO v2 space

README.md

@@ -0,0 +1,12 @@
+---
+title: FSI_ECHO
+emoji: 🧬
+colorFrom: purple
+colorTo: indigo
+sdk: gradio
+sdk_version: "5.0.0"
+app_file: app.py
+pinned: false
+license: apache-2.0
+---
+# FSI_ECHO - 2.6M Param Code AI

app.py

@@ -0,0 +1,51 @@
+import sys, os, json, torch, requests, time
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from fsi_echo import FSIEchoModel, CodeTokenizer, ClosedLoopDebugger
+import gradio as gr
+
+@torch.no_grad()
+def load():
+    MODEL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'prod2_final.pt')
+    if not os.path.exists(MODEL_PATH):
+        url = "https://huggingface.co/FerrellSyntheticIntelligence/FSI_ECHO/resolve/main/prod2_final.pt"
+        r = requests.get(url, stream=True, timeout=300)
+        r.raise_for_status()
+        with open(MODEL_PATH, 'wb') as f:
+            for chunk in r.iter_content(chunk_size=8192):
+                f.write(chunk)
+    ckpt = torch.load(MODEL_PATH, map_location='cpu', weights_only=True)
+    m = FSIEchoModel()
+    m.load_state_dict(ckpt['model'])
+    m.eval()
+    return m, CodeTokenizer()
+
+model, tok = load()
+debugger = ClosedLoopDebugger(model, tok)
+
+def generate(prompt, temp, max_t):
+    t0 = time.time()
+    r = model.generate(tok, prompt, max_tokens=int(max_t), temperature=float(temp), top_k=5, top_p=0.95)
+    return r['generated'], f"{r['tokens']}t | conf:{r['confidence']:.2f} | {time.time()-t0:.1f}s"
+
+def debug(code):
+    r = debugger.debug(code)
+    return r.get('code', 'Could not fix')
+
+with gr.Blocks(title="FSI_ECHO", theme=gr.themes.Monochrome()) as demo:
+    gr.Markdown("# 🧬 FSI_ECHO — 2.6M Param Code AI\n*Morphing Code Swarm*")
+    with gr.Tab("Generate"):
+        p = gr.Textbox(label="Prompt", value="def is_even")
+        t = gr.Slider(0.1, 1.0, 0.1, label="Temperature")
+        m = gr.Slider(10, 200, 50, label="Max Tokens")
+        b = gr.Button("Generate")
+        o = gr.Textbox(label="Output", lines=8)
+        s = gr.Textbox(label="Stats")
+        b.click(generate, [p, t, m], [o, s])
+    with gr.Tab("Debug"):
+        c = gr.Textbox(label="Buggy Code", value="def add(a, b):\n    a + b", lines=6)
+        b2 = gr.Button("Debug")
+        o2 = gr.Textbox(label="Fixed Code", lines=8)
+        b2.click(debug, [c], [o2])
+
+if __name__ == '__main__':
+    demo.launch(server_name='0.0.0.0', server_port=7860)

fsi_echo.py

@@ -0,0 +1,486 @@
+#!/usr/bin/env python3
+"""
+FSI_ECHO - Morphing Code Swarm
+Novel architecture: token morph embedding + nanobot swarm + assembly blocks + self-verification.
+2.6M params — fits in 1.3MB at q4, runs on any phone.
+"""
+import os, sys, json, time, math, random, re, struct
+from typing import List, Dict, Optional, Tuple
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# =============================================================================
+# 1. TOKENIZER
+# =============================================================================
+class CodeTokenizer:
+    SPECIAL = {
+        '<PAD>': 0, '<EOS>': 1, '<BOS>': 2, '<UNK>': 3,
+        '<BUG>': 4, '<FIX>': 5, '<CODE>': 6, '<EXPLAIN>': 7,
+        '<MORPH>': 8, '<ASSEMBLE>': 9, '<SCOUT>': 10, '<COMBAT>': 11,
+    }
+    def __init__(self, vocab_size: int = 4096):
+        self.vocab_size = vocab_size
+        self.vocab = dict(self.SPECIAL)
+        self.inverse = {v: k for k, v in self.SPECIAL.items()}
+        self.next_id = len(self.SPECIAL)
+        self._build()
+    def _build(self):
+        for i in range(32, 127):
+            self._add(chr(i))
+        for t in ['def','class','return','if','else','elif','for','while','in','not',
+            'and','or','import','from','as','try','except','finally','raise','with',
+            'pass','break','continue','yield','lambda','self','None','True','False',
+            'async','await','global','nonlocal','assert','del','print','len','range',
+            'int','str','float','list','dict','set','tuple','type','is','isinstance',
+            'hasattr','getattr','setattr','super','open','Exception','ValueError',
+            'TypeError','KeyError','IndexError','AttributeError','ImportError',
+            'Error','Warning','property','staticmethod','classmethod']:
+            self._add(t)
+        for s in ['==','!=','<=','>=','->','+=','-=','*=','/=','//=','**=','%=',
+                  '<<','>>','**','//','::','=>','++','--','...']:
+            self._add(s)
+        for t in ['fn','func','function','const','let','var','this','typeof','void',
+            'null','undefined','prototype','module','exports','require','new','delete',
+            'throw','catch','switch','case','default','do','while','interface','enum',
+            'implements','private','public','protected','abstract','final','static',
+            'package','boolean','byte','char','double','float','int','long','short',
+            'printf','scanf','malloc','free','sizeof','typedef','struct','union',
+            'include','define','template','typename','namespace','using','virtual',
+            'override','friend','operator','inline','explicit','string','vector',
+            'map','set','auto','decltype','noexcept','constexpr','std','cout','cin',
+            'endl','printf','scanf','NULL','nullptr','true','false','bool']:
+            self._add(t)
+        while self.next_id < self.vocab_size:
+            self._add(f'v{self.next_id}')
+    def _add(self, t):
+        if t not in self.vocab and self.next_id < self.vocab_size:
+            self.vocab[t] = self.next_id
+            self.inverse[self.next_id] = t
+            self.next_id += 1
+    def encode(self, text: str, bos: bool = True, eos: bool = False) -> List[int]:
+        ids = []
+        if bos:
+            ids.append(2)
+        for token in re.findall(r'<[^>]+>|[A-Za-z_][A-Za-z0-9_]*|\.\.\.|==|!=|<=|>=|->|\*\*|//|::|=>|\d+\.\d*|\d+|\S', text):
+            if token in self.vocab:
+                ids.append(self.vocab[token])
+            elif token.lower() in self.vocab:
+                ids.append(self.vocab[token.lower()])
+            else:
+                for ch in token:
+                    if ch in self.vocab:
+                        ids.append(self.vocab[ch])
+                    else:
+                        ids.append(3)
+        if eos:
+            ids.append(1)
+        return ids[:2048]
+    def decode(self, ids: List[int], skip_special: bool = True) -> str:
+        tokens = []
+        for i in ids:
+            if i in self.inverse:
+                t = self.inverse[i]
+                if skip_special and t.startswith('<') and t.endswith('>'):
+                    continue
+                tokens.append(t)
+            else:
+                tokens.append(' ')
+        return ''.join(tokens)
+    @property
+    def pad_id(self): return 0
+    @property
+    def eos_id(self): return 1
+    @property
+    def bos_id(self): return 2
+    @property
+    def vocab_size_(self): return len(self.vocab)
+
+# =============================================================================
+# 2. MORPH EMBEDDING
+# =============================================================================
+class MorphEmbedding(nn.Module):
+    def __init__(self, vocab_size: int, d_model: int, morph_width: int = 3):
+        super().__init__()
+        self.d_model = d_model
+        self.morph_width = morph_width
+        self.base_embed = nn.Embedding(vocab_size, d_model)
+        # Causal morph: only current + past (morph_width tokens, no future leakage)
+        self.morph_net = nn.Sequential(
+            nn.Linear(d_model * morph_width, d_model), nn.Tanh(),
+            nn.Linear(d_model, d_model),
+        )
+        self.gate = nn.Linear(d_model * 2, d_model)
+
+    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
+        B, T = tokens.shape
+        base = self.base_embed(tokens)
+        # Causal sliding window: each position only sees itself and preceding tokens
+        padded = F.pad(base, (0, 0, self.morph_width - 1, 0), mode='replicate')
+        contexts = []
+        for i in range(self.morph_width):
+            contexts.append(padded[:, i:i+T, :])
+        context = torch.cat(contexts, dim=-1)
+        morph = self.morph_net(context)
+        gate = torch.sigmoid(self.gate(torch.cat([base, morph], dim=-1)))
+        return base + gate * morph
+
+# =============================================================================
+# 3. NANOBOT SWARM
+# =============================================================================
+class NanobotSwarm(nn.Module):
+    def __init__(self, d_model: int, n_nanobots: int = 512, scout_dim: int = 64, combat_dim: int = 128):
+        super().__init__()
+        self.d_model = d_model
+        self.n_nanobots = n_nanobots
+        self.nano_keys = nn.Parameter(torch.randn(n_nanobots, d_model // 4))
+        self.nano_vals = nn.Parameter(torch.randn(n_nanobots, d_model))
+        self.scout_ffn = nn.Sequential(
+            nn.Linear(d_model, scout_dim), nn.GELU(), nn.Linear(scout_dim, d_model),
+        )
+        self.combat_ffn = nn.Sequential(
+            nn.Linear(d_model, combat_dim), nn.GELU(), nn.Linear(combat_dim, d_model),
+        )
+        self.mode_router = nn.Linear(d_model, 2)
+        self.assembly = nn.Linear(d_model, d_model)
+        self.norm = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(0.1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, D = x.shape
+        residual = x
+        x = self.norm(x)
+        router_query = x[..., :self.d_model // 4]
+        nano_scores = torch.matmul(router_query, self.nano_keys.T)
+        nano_weights = F.softmax(nano_scores / math.sqrt(self.d_model // 4), dim=-1)
+        mode_logits = self.mode_router(x)
+        mode_w = F.softmax(mode_logits, dim=-1)
+        scout_out = self.scout_ffn(x)
+        combat_out = self.combat_ffn(x)
+        mode_out = mode_w[:, :, 0:1] * scout_out + mode_w[:, :, 1:2] * combat_out
+        nano_out = torch.matmul(nano_weights, self.nano_vals)
+        out = residual + self.dropout(self.assembly(mode_out + nano_out))
+        return out
+
+# =============================================================================
+# 4. ASSEMBLY BLOCK
+# =============================================================================
+class AssemblyBlock(nn.Module):
+    def __init__(self, d_model: int, n_heads: int = 4, dropout: float = 0.1):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
+        self.proj = nn.Linear(d_model, d_model, bias=False)
+        self.ffn = nn.Sequential(
+            nn.Linear(d_model, d_model * 2), nn.GELU(),
+            nn.Linear(d_model * 2, d_model), nn.Dropout(dropout),
+        )
+        self.adapt_gate = nn.Linear(d_model, 1)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, D = x.shape
+        qkv = self.qkv(self.norm1(x))
+        qkv = qkv.reshape(B, T, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+        scale = math.sqrt(self.head_dim)
+        attn = torch.matmul(q, k.transpose(-2, -1)) / scale
+        mask = torch.triu(torch.ones(T, T, device=x.device) * float('-inf'), diagonal=1)
+        attn = attn + mask.unsqueeze(0)
+        attn = F.softmax(attn, dim=-1)
+        out = torch.matmul(attn, v).transpose(1, 2).reshape(B, T, D)
+        out = self.proj(out)
+        x = x + self.dropout(out)
+        gate = torch.sigmoid(self.adapt_gate(self.norm2(x)))
+        x = x + gate * self.dropout(self.ffn(self.norm2(x)))
+        return x
+
+# =============================================================================
+# 5. FSI_ECHO MODEL
+# =============================================================================
+class FSIEchoModel(nn.Module):
+    def __init__(self, vocab_size: int = 4096, d_model: int = 192,
+                 n_swarm_layers: int = 3, n_assembly_layers: int = 3,
+                 n_nanobots: int = 512):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.morph_embed = MorphEmbedding(vocab_size, d_model)
+        self.swarm_layers = nn.ModuleList([
+            NanobotSwarm(d_model, n_nanobots) for _ in range(n_swarm_layers)
+        ])
+        self.assembly_layers = nn.ModuleList([
+            AssemblyBlock(d_model) for _ in range(n_assembly_layers)
+        ])
+        self.norm = nn.LayerNorm(d_model)
+        self.verify = nn.Sequential(
+            nn.Linear(d_model, d_model // 2), nn.GELU(), nn.Linear(d_model // 2, 1),
+        )
+        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
+        self.lm_head.weight = self.morph_embed.base_embed.weight
+        self._init_weights()
+
+    def _init_weights(self):
+        for m in self.modules():
+            if isinstance(m, (nn.Linear, nn.Embedding)):
+                nn.init.normal_(m.weight, mean=0.0, std=0.02)
+                if hasattr(m, 'bias') and m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.LayerNorm):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+
+    def forward(self, tokens: torch.Tensor, targets: Optional[torch.Tensor] = None) -> Dict:
+        x = self.morph_embed(tokens)
+        for layer in self.swarm_layers:
+            x = layer(x)
+        for layer in self.assembly_layers:
+            x = layer(x)
+        x = self.norm(x)
+        logits = self.lm_head(x)
+        confidence = torch.sigmoid(self.verify(x)).squeeze(-1)
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, self.vocab_size), targets.view(-1), ignore_index=0)
+        return {'logits': logits, 'confidence': confidence, 'loss': loss}
+
+    @torch.no_grad()
+    def generate(self, tokenizer, prompt: str, max_tokens: int = 512,
+                 temperature: float = 0.3, top_k: int = 5, top_p: float = 0.9) -> Dict:
+        self.eval()
+        device = next(self.parameters()).device
+        toks = torch.tensor([tokenizer.encode(prompt)], dtype=torch.long, device=device)
+        generated_ids = []
+        confs = []
+        for _ in range(min(max_tokens, 2048 - toks.shape[1])):
+            out = self.forward(toks)
+            logits = out['logits'][0, -1, :] / max(temperature, 0.01)
+            if top_k > 0:
+                vals, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < vals[-1]] = float('-inf')
+            if top_p < 1.0:
+                sorted_lg, sorted_idx = torch.sort(logits, descending=True)
+                cum = torch.cumsum(F.softmax(sorted_lg, dim=-1), dim=-1)
+                rm = cum > top_p
+                rm[1:] = rm[:-1].clone()
+                rm[0] = False
+                logits[sorted_idx[rm]] = float('-inf')
+            logits = torch.nan_to_num(logits, nan=-100.0, posinf=100.0, neginf=-100.0)
+            logits[0] = float('-inf')
+            for rid in range(300, logits.size(-1)):
+                logits[rid] = float('-inf')
+            if (logits > float(-1e9)).sum() == 0:
+                logits[tokenizer.eos_id] = 0.0
+            nxt = logits.argmax().unsqueeze(0)
+            confs.append(out['confidence'][0, -1].item())
+            if nxt.item() == tokenizer.eos_id:
+                break
+            generated_ids.append(nxt.item())
+            toks = torch.cat([toks, nxt.unsqueeze(0)], dim=1)
+        generated = tokenizer.decode(generated_ids, skip_special=True)
+        avg_conf = sum(confs) / max(len(confs), 1)
+        return {'generated': generated, 'confidence': avg_conf, 'tokens': len(generated_ids)}
+
+    def param_count(self) -> int:
+        return sum(p.numel() for p in self.parameters())
+
+# =============================================================================
+# 6. TRAINER
+# =============================================================================
+class Trainer:
+    def __init__(self, model: FSIEchoModel, tokenizer: CodeTokenizer, lr: float = 3e-4):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.optimizer = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.1)
+        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=5000, eta_min=1e-5)
+
+    def step(self, texts: List[str], batch_size: int = 2, device: str = 'cpu') -> float:
+        self.model.train()
+        batch = random.sample(texts, min(batch_size, len(texts)))
+        encoded = []
+        max_len = 0
+        for t in batch:
+            toks = self.tokenizer.encode(t)
+            if 5 < len(toks) <= 2048:
+                encoded.append(toks)
+                max_len = max(max_len, len(toks))
+        if not encoded:
+            return 0.0
+        padded = torch.zeros(len(encoded), max_len, dtype=torch.long)
+        targets = torch.full((len(encoded), max_len), 0, dtype=torch.long)
+        for i, toks in enumerate(encoded):
+            padded[i, :len(toks)] = torch.tensor(toks)
+            targets[i, :len(toks)-1] = torch.tensor(toks[1:])
+            targets[i, len(toks)-1] = 1
+        padded, targets = padded.to(device), targets.to(device)
+        out = self.model(padded, targets)
+        self.optimizer.zero_grad()
+        out['loss'].backward()
+        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
+        self.optimizer.step()
+        self.scheduler.step()
+        return out['loss'].item()
+
+    def save(self, path: str):
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        torch.save({
+            'model': self.model.state_dict(),
+            'optimizer': self.optimizer.state_dict(),
+            'config': {'vocab_size': self.model.vocab_size, 'd_model': self.model.d_model,
+                       'n_swarm': len(self.model.swarm_layers), 'n_assembly': len(self.model.assembly_layers),
+                       'n_nanobots': self.model.swarm_layers[0].n_nanobots if self.model.swarm_layers else 512},
+        }, path)
+
+    @classmethod
+    def load(cls, path: str, device: str = 'cpu') -> 'Trainer':
+        data = torch.load(path, map_location=device, weights_only=True)
+        cfg = data.get('config', {})
+        model = FSIEchoModel(
+            vocab_size=cfg.get('vocab_size', 4096), d_model=cfg.get('d_model', 192),
+            n_swarm_layers=cfg.get('n_swarm', 3), n_assembly_layers=cfg.get('n_assembly', 3),
+            n_nanobots=cfg.get('n_nanobots', 512),
+        )
+        model.load_state_dict(data['model'])
+        tokenizer = CodeTokenizer(vocab_size=cfg.get('vocab_size', 4096))
+        t = cls(model, tokenizer)
+        if 'optimizer' in data:
+            t.optimizer.load_state_dict(data['optimizer'])
+        model.to(device)
+        return t
+
+# =============================================================================
+# 7. CLOSED-LOOP DEBUG
+# =============================================================================
+class CodeVerifier:
+    @staticmethod
+    def check_syntax(code: str) -> Tuple[bool, str]:
+        try:
+            compile(code, '<debug>', 'exec')
+            return True, ""
+        except SyntaxError as e:
+            return False, f"Line {e.lineno}: {e.msg}"
+    @staticmethod
+    def find_issues(code: str) -> List[str]:
+        issues = []
+        for i, line in enumerate(code.split('\n'), 1):
+            s = line.strip()
+            if s == 'except:':
+                issues.append(f"L{i}: Bare except — specify exception")
+        return issues
+    def verify(self, code: str) -> Dict:
+        ok, err = self.check_syntax(code)
+        issues = self.find_issues(code)
+        return {'valid': ok and not issues, 'syntax_ok': ok, 'error': err, 'issues': issues, 'code': code}
+
+class ClosedLoopDebugger:
+    def __init__(self, model: FSIEchoModel, tokenizer: CodeTokenizer, max_iters: int = 3):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.max_iters = max_iters
+        self.verifier = CodeVerifier()
+    def debug(self, code: str, requirement: str = "", max_iterations: int = None) -> Dict:
+        iters = max_iterations or self.max_iters
+        result = self._extract_code(code)
+        if not result:
+            return {'code': code, 'error': 'Could not parse code', 'iterations': 0, 'confidence': 0.0}
+        buggy = result
+        prompt = f"Fix this code:\n```python\n{buggy}\n```\nFixed:\n```python\n"
+        best_code, best_v = buggy, self.verifier.verify(buggy)
+        for i in range(iters):
+            gen = self.model.generate(self.tokenizer, prompt, max_tokens=256, temperature=0.3, top_k=5)
+            fixed = self._extract_code(gen['generated'])
+            if not fixed:
+                prompt += "\n```python\n"
+                continue
+            v = self.verifier.verify(fixed)
+            if v['valid']:
+                return {'code': fixed, 'verification': v, 'iterations': i+1, 'confidence': gen['confidence']}
+            if len(v['issues']) < len(best_v['issues']):
+                best_code, best_v = fixed, v
+            if v['issues']:
+                prompt += f"\nIssues: {'; '.join(v['issues'])}\nFixed:\n```python\n"
+            elif not v['syntax_ok']:
+                prompt += f"\n{v['error']}\nFixed:\n```python\n"
+            else:
+                break
+        return {'code': best_code, 'verification': best_v, 'iterations': iters, 'confidence': 0.0}
+    def _extract_code(self, text: str) -> str:
+        m = re.search(r'```(?:python)?\n(.*?)```', text, re.DOTALL)
+        if m: return m.group(1).strip()
+        lines = text.split('\n')
+        code = []
+        in_code = False
+        for line in lines:
+            if line.startswith('```'): in_code = not in_code; continue
+            if in_code: code.append(line)
+        return '\n'.join(code).strip() if code else ''
+
+# =============================================================================
+# 8. GGUF EXPORT
+# =============================================================================
+def export_gguf(model: FSIEchoModel, tokenizer: CodeTokenizer, path: str):
+    sd = model.state_dict()
+    keys = sorted(sd.keys())
+    meta = {
+        'general.name': 'FSI_ECHO', 'general.architecture': 'fsi_echo',
+        'general.description': 'Morphing Code Swarm',
+        'general.file_type': 0, 'general.vocab_size': model.vocab_size,
+        'general.context_length': 2048, 'general.parameter_count': model.param_count(),
+        'fsi_echo.block_count': len(model.swarm_layers) + len(model.assembly_layers),
+        'fsi_echo.embedding_length': model.d_model,
+        'fsi_echo.feed_forward_length': model.d_model * 2,
+        'fsi_echo.attention.head_count': 4,
+        'fsi_echo.nanobot_count': model.swarm_layers[0].n_nanobots if model.swarm_layers else 512,
+    }
+    with open(path, 'wb') as f:
+        f.write(b'GGUF' + struct.pack('<I', 3) + struct.pack('<Q', len(keys)) + struct.pack('<Q', len(meta)))
+        for k, v in meta.items():
+            f.write(struct.pack('<I', len(k)) + k.encode())
+            if isinstance(v, str):
+                f.write(struct.pack('<I', 8) + struct.pack('<Q', len(v)) + v.encode())
+            elif isinstance(v, int):
+                f.write(struct.pack('<I', 4) + struct.pack('<I', v))
+            elif isinstance(v, float):
+                f.write(struct.pack('<I', 6) + struct.pack('<f', v))
+        offset = 0
+        for name in keys:
+            nb = len(name)
+            f.write(struct.pack('<Q', nb) + name.encode() +
+                    struct.pack('<I', len(sd[name].shape)) +
+                    b''.join(struct.pack('<Q', d) for d in sd[name].shape) +
+                    struct.pack('<I', 0) + struct.pack('<Q', offset))
+            offset += sd[name].numel() * 4
+        for name in keys:
+            f.write(sd[name].float().numpy().tobytes())
+
+# =============================================================================
+# 9. MAIN
+# =============================================================================
+if __name__ == '__main__':
+    import argparse
+    p = argparse.ArgumentParser(description='FSI_ECHO')
+    p.add_argument('--train', action='store_true')
+    p.add_argument('--gen', type=str, help='Generate from prompt')
+    p.add_argument('--gguf', type=str, help='Export to GGUF')
+    p.add_argument('--load', type=str, help='Load checkpoint')
+    p.add_argument('--steps', type=int, default=5000)
+    args = p.parse_args()
+    device = 'cpu'
+    if args.load and os.path.exists(args.load):
+        trainer = Trainer.load(args.load, device)
+    else:
+        model = FSIEchoModel()
+        model.to(device)
+        tok = CodeTokenizer()
+        print(f"New model: {model.param_count():,} params")
+        trainer = Trainer(model, tok)
+    if args.gen:
+        r = trainer.model.generate(trainer.tokenizer, args.gen, max_tokens=256)
+        print(r['generated'])
+    if args.gguf:
+        export_gguf(trainer.model, trainer.tokenizer, args.gguf)

requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.0.0
+gradio>=5.0.0
+requests>=2.0.0