"""Transformers+peft backend for the real Codette orchestrator, for ZeroGPU. The real CodetteOrchestrator is built on llama.cpp (GGUF). llama.cpp can't use ZeroGPU's per-request GPU (that works through a torch hijack). So we keep the ENTIRE faithful pipeline — routing, the behavioral locks, the integrity/complexity/ role layer, constraint override, chaos mitigation, self-correction, synthesis — and swap ONLY the model layer: * self._llm.create_chat_completion(...) -> a transformers+peft shim * self._load_model(adapter) (llama.cpp hot-swap) -> peft set_adapter Everything above the primitive is reused unchanged from codette_orchestrator.py. """ import os import sys import types from pathlib import Path def _install_llama_cpp_stub(): """codette_orchestrator imports llama_cpp at module load. We never call the real GGUF backend, so a stub satisfies the import without installing it.""" if "llama_cpp" in sys.modules: return m = types.ModuleType("llama_cpp") m.Llama = object m.llama_adapter_lora_init = lambda *a, **k: None m.llama_set_adapter_lora = lambda *a, **k: 0 m.llama_set_adapters_lora = lambda *a, **k: 0 sys.modules["llama_cpp"] = m # ---- Generation shim ----------------------------------------------------- class _TransformersLLM: """Mimics the slice of llama-cpp-python's API that the orchestrator uses: create_chat_completion(messages, max_tokens, temperature, top_p, repeat_penalty, stop, ...) -> OpenAI-style dict.""" def __init__(self, model, tokenizer, mock=False): self.model = model self.tokenizer = tokenizer self.mock = mock def create_chat_completion(self, messages, max_tokens=512, temperature=0.7, top_p=0.9, repeat_penalty=1.1, stop=None, **_): if self.mock: text = f"[mock reply to: {messages[-1]['content'][:60]}]" return {"choices": [{"message": {"content": text}}], "usage": {"completion_tokens": len(text.split())}} import torch tok = self.tokenizer enc = tok.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt", return_dict=True, ).to(self.model.device) # Llama-3.1-Instruct ends assistant turns with <|eot_id|>, NOT the # tokenizer's default eos (<|end_of_text|>). Without this the model # sails past its own stop token and rambles to max_tokens. eos_ids = [tok.eos_token_id] eot = tok.convert_tokens_to_ids("<|eot_id|>") if isinstance(eot, int) and eot >= 0 and eot not in eos_ids: eos_ids.append(eot) gen = self.model.generate( **enc, max_new_tokens=min(int(max_tokens), 512), do_sample=temperature > 0, temperature=max(temperature, 0.01), top_p=top_p, repetition_penalty=max(repeat_penalty, 1.15), no_repeat_ngram_size=4, eos_token_id=eos_ids, pad_token_id=tok.pad_token_id, ) new = gen[0][enc["input_ids"].shape[1]:] text = self.tokenizer.decode(new, skip_special_tokens=True).strip() for s in (stop or []): i = text.find(s) if i != -1: text = text[:i] return {"choices": [{"message": {"content": text.strip()}}], "usage": {"completion_tokens": int(new.shape[0])}} def build_orchestrator(base_model, adapters, *, mock=False, verbose=True, device="cuda", dtype=None): """Construct a transformers-backed CodetteOrchestrator subclass. adapters: {name: local_path_to_peft_adapter_dir} mock: skip real model load; shim returns canned text (for wiring tests). """ _install_llama_cpp_stub() import codette_orchestrator as _co from codette_orchestrator import CodetteOrchestrator, AdapterRouter # Non-invasive shim: codette_orchestrator.generate() references bare `re` # at module scope but only imported `re as _re_mod`. Inject it rather than # editing the source. (Real fix belongs in the repo: add `import re`.) if not hasattr(_co, "re"): import re as _re _co.re = _re class HFCodetteOrchestrator(CodetteOrchestrator): def __init__(self): # --- replicate the parts of the parent __init__ we need, minus the # GGUF hot-swap init --- self.n_ctx = 8192 self.n_gpu_layers = 0 self.verbose = verbose self.memory_weighting = None self._llm = None self._current_adapter = "__unset__" self._adapter_handles = {} self._model_ptr = None self._ctx_ptr = None self._hf_adapters = dict(adapters) self.available_adapters = list(adapters.keys()) self.router = AdapterRouter( available_adapters=self.available_adapters, memory_weighting=None, ) print(f"Available adapters: {', '.join(self.available_adapters)}") # integrity layer (same try/except as parent) try: from reasoning_forge.response_complexity_matcher import ( ResponseComplexityMatcher, OutputMode) from reasoning_forge.conversation_role_tracker import ( ConversationRoleTracker) self._complexity_matcher = ResponseComplexityMatcher() self._role_tracker = ConversationRoleTracker() self._OutputMode = OutputMode print(" Integrity layer loaded") except Exception as e: print(f" NOTE: integrity layer not loaded: {e}") self._complexity_matcher = None self._role_tracker = None self._OutputMode = None try: from reasoning_forge.style_adaptive_synthesis import ( StyleAdaptiveSynthesis) self._style_adapter = StyleAdaptiveSynthesis() print(" Style layer loaded") except Exception as e: print(f" NOTE: style layer not loaded: {e}") self._style_adapter = None self._hf_load() def _hf_load(self): if mock: self.model = None self.tokenizer = _MockTok() self._llm = _TransformersLLM(None, self.tokenizer, mock=True) print(" [mock] transformers backend stubbed") return import torch from transformers import AutoModelForCausalLM, AutoTokenizer from peft import PeftModel d = dtype or torch.bfloat16 print(f" loading base {base_model} …") self.tokenizer = AutoTokenizer.from_pretrained(base_model) if self.tokenizer.pad_token_id is None: self.tokenizer.pad_token_id = self.tokenizer.eos_token_id base = AutoModelForCausalLM.from_pretrained(base_model, dtype=d) names = list(self._hf_adapters) _real = torch.cuda.is_available torch.cuda.is_available = lambda: False # keep peft off CUDA at load try: m = PeftModel.from_pretrained( base, self._hf_adapters[names[0]], adapter_name=names[0]) for n in names[1:]: m.load_adapter(self._hf_adapters[n], adapter_name=n) finally: torch.cuda.is_available = _real m.eval() if device == "cuda": m.to("cuda") self.model = m self._llm = _TransformersLLM(m, self.tokenizer, mock=False) print(" transformers backend ready") def _load_model(self, adapter_name=None): if adapter_name == self._current_adapter: return if not mock and self.model is not None: if adapter_name and adapter_name in self._hf_adapters: self.model.set_adapter(adapter_name) self.model.enable_adapter_layers() else: self.model.disable_adapter_layers() # base model self._current_adapter = adapter_name if self.verbose: print(f" [swapped to {adapter_name or 'base'}]") # neutralize the GGUF hot-swap init the parent would run def _init_hotswap(self): pass return HFCodetteOrchestrator() class _MockTok: pad_token_id = 0 eos_token_id = 0 def apply_chat_template(self, *a, **k): return ""