Create evo_plugin.py

evo_plugin.py

@@ -0,0 +1,78 @@
+# evo_plugin.py — REAL Evo decoder integration (state_dict -> generation)
+import os, torch
+from typing import Optional
+
+# If you use SentencePiece:
+try:
+    import sentencepiece as spm
+    _HAS_SPM = True
+except Exception:
+    _HAS_SPM = False
+
+from evo_core_gpt import EvoGPT
+
+MODEL_PATH  = os.environ.get("EVO_DECODER_PATH", "models/evo_decoder.pt")
+SPM_PATH    = os.environ.get("EVO_SPM_PATH", "models/evo_tokenizer.model")   # SentencePiece file
+NUM_HEADS   = int(os.environ.get("EVO_NUM_HEADS", "8"))                      # <-- set this to your trained value
+
+class _SPTokenizer:
+    def __init__(self, spm_path: str):
+        if not _HAS_SPM:
+            raise RuntimeError("sentencepiece not installed; add 'sentencepiece' to requirements.txt")
+        self.sp = spm.SentencePieceProcessor(model_file=spm_path)
+        # Try to detect special tokens if your model has them in the SPM model; else leave None
+        self.bos_id = self.sp.bos_id() if self.sp.bos_id() >= 0 else None
+        self.eos_id = self.sp.eos_id() if self.sp.eos_id() >= 0 else None
+
+    def encode(self, text: str):
+        return self.sp.encode(text, out_type=int)
+
+    def decode(self, ids):
+        return self.sp.decode(ids)
+
+class EvoTextGenerator:
+    def __init__(self, weights_path: str = MODEL_PATH, spm_path: str = SPM_PATH, num_heads: int = NUM_HEADS, device: str = "cpu"):
+        self.device = torch.device(device)
+        # 1) peek shapes from state_dict
+        sd = torch.load(weights_path, map_location="cpu")
+        # infer config
+        vocab_size, d_model = sd["token_emb.weight"].shape
+        n_positions = sd["pos_emb.weight"].shape[0]
+        # count layers by scanning keys "blocks.N."
+        import re
+        layer_ids = sorted({int(re.findall(r"blocks\.(\d+)\.", k)[0]) for k in sd.keys() if k.startswith("blocks.")})
+        n_layers = 1 + max(layer_ids) if layer_ids else 1
+
+        # 2) build model and load weights
+        self.model = EvoGPT(vocab_size, d_model, n_layers, n_positions, num_heads=num_heads).to(self.device)
+        missing, unexpected = self.model.load_state_dict(sd, strict=False)
+        # Usually both lists should be empty; print if needed:
+        if missing or unexpected:
+            print("[EvoGPT] load_state_dict notice -> missing:", missing, "unexpected:", unexpected)
+        self.model.eval()
+
+        # 3) tokenizer
+        if not os.path.exists(spm_path):
+            raise FileNotFoundError(f"Tokenizer not found at {spm_path}. Provide SentencePiece model via models/evo_tokenizer.model or set EVO_SPM_PATH.")
+        self.tok = _SPTokenizer(spm_path)
+
+    @torch.no_grad()
+    def generate(self, prompt: str, max_new_tokens: int = 200, temperature: float = 0.4) -> str:
+        # encode
+        ids = self.tok.encode(prompt)
+        if self.tok.bos_id is not None and (not ids or ids[0] != self.tok.bos_id):
+            ids = [self.tok.bos_id] + ids
+        input_ids = torch.tensor([ids], dtype=torch.long, device=self.device)
+        # generate
+        out_ids = self.model.generate(
+            input_ids,
+            max_new_tokens=int(max_new_tokens),
+            temperature=float(temperature),
+            eos_token_id=self.tok.eos_id,
+        )[0].tolist()
+        # strip the prompt tokens if desired; for now decode all
+        text = self.tok.decode(out_ids)
+        return text
+
+def load_model():
+    return EvoTextGenerator()