Add controlled inference engine

nomen_ai/inference.py

@@ -0,0 +1,44 @@
+"""Controlled generation engine: control vector + creativity decoding + anti-dup filter."""
+from typing import List
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from .control import ControlVector, SYSTEM_PROMPT
+from .antidup import AntiDuplicationMatrix
+from .phonetics import count_syllables, char_len
+
+def creativity_to_decoding(c: float) -> dict:
+    if c <= 0.25:
+        return dict(do_sample=False, penalty_alpha=0.6, top_k=5, repetition_penalty=1.2, no_repeat_ngram_size=2)
+    return dict(do_sample=True, temperature=round(0.7+c*1.3,2), min_p=round(max(0.02,0.12-c*0.08),3), repetition_penalty=1.25, no_repeat_ngram_size=2)
+
+class NomenAI:
+    def __init__(self, model_id, base_model=None, device=None, antidup=None):
+        self.device=device or ('cuda' if torch.cuda.is_available() else 'cpu')
+        self.tok=AutoTokenizer.from_pretrained(model_id if not base_model else base_model)
+        if base_model:
+            from peft import PeftModel
+            base=AutoModelForCausalLM.from_pretrained(base_model, torch_dtype=torch.float16).to(self.device)
+            self.model=PeftModel.from_pretrained(base, model_id).to(self.device)
+        else:
+            self.model=AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16).to(self.device)
+        self.model.eval(); self.antidup=antidup or AntiDuplicationMatrix()
+    def _prompt_ids(self, cv: ControlVector):
+        msgs=[{'role':'system','content':SYSTEM_PROMPT},{'role':'user','content':cv.to_prompt()}]
+        text=self.tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+        return self.tok(text, return_tensors='pt').to(self.device)
+    @torch.no_grad()
+    def generate(self, cv: ControlVector, n=10, tol_syl=1, tol_len=3, enforce=True, max_new_tokens=12) -> List[dict]:
+        cv.validate(); enc=self._prompt_ids(cv); dec=creativity_to_decoding(cv.creativity)
+        out=self.model.generate(**enc, num_return_sequences=n*3, max_new_tokens=max_new_tokens, pad_token_id=self.tok.eos_token_id, **dec)
+        gen=out[:, enc['input_ids'].shape[1]:]
+        results=[]; seen=set()
+        for g in gen:
+            txt=self.tok.decode(g, skip_special_tokens=True).strip(); name=''.join(ch for ch in (txt.split()[0] if txt else '') if ch.isalpha())
+            if not name or name.lower() in seen: continue
+            nov=self.antidup.novelty(name); syl=count_syllables(name); clen=char_len(name)
+            ok=self.antidup.is_novel(name)
+            if enforce: ok=ok and abs(syl-cv.syllables)<=tol_syl and abs(clen-cv.char_len)<=tol_len
+            if ok:
+                seen.add(name.lower()); results.append({'name':name.capitalize(),'novelty':nov,'syllables':syl,'chars':clen})
+            if len(results)>=n: break
+        return sorted(results, key=lambda r:-r['novelty'])