Spaces:
Running on Zero
Running on Zero
| """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 "" | |