"""VocabFusion-Pythia160m-Fused — OpenAI-compatible inference server. VocabFusion-SymbioGPT architecture (donor embedding + junction + SymbioGPT body) with knowledge from Pythia-160M (12 layers, d=768) projected in via PCA decomposition and blended at alpha=0.3. 17.5M params. Downloads checkpoint from model repo and tokenizer from VocabFusion-SymbioGPT-10M. """ import json as json_mod import math import os import time import uuid import torch import torch.nn.functional as F import uvicorn from fastapi import FastAPI, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse, StreamingResponse from huggingface_hub import hf_hub_download from tokenizers import Tokenizer from symbio_model import SymbioConfig from vocab_fusion_model import VocabFusionSymbioGPT # ═══════════════════════════════════════════════════════════════════ # Configuration # ═══════════════════════════════════════════════════════════════════ MODEL_REPO = os.environ.get("MODEL_REPO", "LisaMegaWatts/VocabFusion-Pythia160m-Fused") TOKENIZER_REPO = os.environ.get("TOKENIZER_REPO", "LisaMegaWatts/VocabFusion-SymbioGPT-10M") PORT = int(os.environ.get("PORT", "7860")) CHECKPOINT_FILE = "vocabfusion_pythia160m_fused.pt" DONOR_VOCAB_SIZE = 50304 DONOR_DIM = 128 MODEL_CONFIG = SymbioConfig( d_model=320, n_layers=8, n_heads=5, head_dim=64, ffn_mult=4, context_length=256, vocab_size=2000, weight_tying=True, organelles=("causal_conv", "monarch", "long_conv", "attention"), conv_kernel_size=4, n_monarch_heads=1, gate_temperature_init=1.0, free_energy_beta=0.001, ) # ═══════════════════════════════════════════════════════════════════ # Load model and tokenizer # ═══════════════════════════════════════════════════════════════════ print(f"Downloading checkpoint from {MODEL_REPO} ...") ckpt_path = hf_hub_download(repo_id=MODEL_REPO, filename=CHECKPOINT_FILE) print(f"Downloading tokenizer from {TOKENIZER_REPO} ...") tok_path = hf_hub_download(repo_id=TOKENIZER_REPO, filename="tokenizer/tokenizer.json") print("Loading tokenizer ...") tokenizer = Tokenizer.from_file(tok_path) actual_vocab_size = tokenizer.get_vocab_size() print(f" Tokenizer vocab_size = {actual_vocab_size}") print(f" Donor embedding size = {DONOR_VOCAB_SIZE} (padded)") print("Loading VocabFusion-Pythia160m-Fused model ...") model = VocabFusionSymbioGPT(MODEL_CONFIG, donor_vocab_size=DONOR_VOCAB_SIZE, donor_dim=DONOR_DIM) state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=True) if "model_state_dict" in state_dict: state_dict = state_dict["model_state_dict"] elif "state_dict" in state_dict: state_dict = state_dict["state_dict"] state_dict = {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()} model.load_state_dict(state_dict) model.eval() n_params = sum(p.numel() for p in model.parameters()) print(f" Model ready: {n_params/1e6:.1f}M params (Pythia-160M knowledge fused)") # ═══════════════════════════════════════════════════════════════════ # Generation # ═══════════════════════════════════════════════════════════════════ @torch.no_grad() def generate_streaming( prompt: str, max_tokens: int = 200, temperature: float = 0.8, top_k: int = 40, top_p: float = 1.0, ): """Generator yielding token strings one at a time for real SSE streaming.""" encoding = tokenizer.encode(prompt) tokens = encoding.ids if not tokens: tokens = [0] idx = torch.tensor([tokens], dtype=torch.long) generated_ids = [] prev_text = "" for _ in range(max_tokens): idx_cond = idx[:, -MODEL_CONFIG.context_length:] logits = model(idx_cond) logits_last = logits[0, -1, :].float() if temperature > 0.01: logits_last = logits_last / temperature else: logits_last = logits_last / 0.01 if 0 < top_k < logits_last.size(0): threshold = torch.topk(logits_last, top_k).values[-1] logits_last[logits_last < threshold] = float("-inf") if top_p < 1.0: sorted_logits, sorted_indices = torch.sort(logits_last, descending=True) probs_sorted = F.softmax(sorted_logits, dim=-1) cumprobs = torch.cumsum(probs_sorted, dim=-1) cutoff_mask = cumprobs - probs_sorted > top_p sorted_logits[cutoff_mask] = float("-inf") logits_last = sorted_logits.scatter(0, sorted_indices, sorted_logits) probs = F.softmax(logits_last, dim=-1) next_id = torch.multinomial(probs, 1).item() if next_id == 0: break idx = torch.cat([idx, torch.tensor([[next_id]])], dim=1) generated_ids.append(next_id) full_text = tokenizer.decode(generated_ids, skip_special_tokens=True) delta = full_text[len(prev_text):] prev_text = full_text if delta: yield delta @torch.no_grad() def generate( prompt: str, max_tokens: int = 200, temperature: float = 0.8, top_k: int = 40, top_p: float = 1.0, ) -> str: """Generate complete text (non-streaming wrapper).""" return "".join(generate_streaming(prompt, max_tokens, temperature, top_k, top_p)) # ═══════════════════════════════════════════════════════════════════ # FastAPI server # ═══════════════════════════════════════════════════════════════════ app = FastAPI() app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) MODEL_CREATED_AT = int(time.time()) MODEL_ID = "vocabfusion-pythia160m-fused" def extract_prompt(messages): if not messages: return "" for msg in reversed(messages): if msg.get("role") == "user": return msg.get("content", "") return messages[-1].get("content", "") @app.get("/") def health(): return { "name": "VocabFusion-Pythia160m-Fused", "version": "1.0.0", "description": ( "VocabFusion-SymbioGPT with knowledge from Pythia-160M (12 layers, d=768) " "projected via PCA decomposition and blended at alpha=0.3. " "The donor embedding (GPTNeoX, ~50k tokens, 128-dim) bridges to " "the 320-dim SymbioGPT body via a junction layer." ), "architecture": "VocabFusion wrapper + 4-organelle decoder " "(CausalConv + Monarch + LongConv + Attention) " "+ OrganelleGate + SwiGLU + RoPE + RMSNorm", "model": { "d_model": MODEL_CONFIG.d_model, "n_layers": MODEL_CONFIG.n_layers, "n_heads": MODEL_CONFIG.n_heads, "context_length": MODEL_CONFIG.context_length, "donor_vocab_size": DONOR_VOCAB_SIZE, "donor_dim": DONOR_DIM, "params": f"{n_params/1e6:.1f}M", }, "fusion": { "source": "EleutherAI/pythia-160m (12 layers, d=768)", "method": "PCA projection + blend", "blend_alpha": 0.3, "pca_avg_variance": 0.921, }, "organelles": list(MODEL_CONFIG.organelles), "endpoints": ["/v1/models", "/v1/chat/completions"], "features": ["streaming", "OpenAI-compatible", "top-k", "top-p"], "github": "https://github.com/DavinciDreams/SymbioGPT", } @app.get("/v1/models") def list_models(): return { "object": "list", "data": [{ "id": MODEL_ID, "object": "model", "created": MODEL_CREATED_AT, "owned_by": "symbiogpt", }], } @app.post("/v1/chat/completions") async def chat_completions(request: Request): try: body = await request.json() except Exception: return JSONResponse(status_code=400, content={ "error": {"message": "Invalid JSON", "type": "invalid_request_error"} }) temperature = max(0.01, min(2.0, body.get("temperature", 0.8))) max_tokens = max(1, min(MODEL_CONFIG.context_length, body.get("max_tokens", 200))) top_k_val = max(0, min(DONOR_VOCAB_SIZE, body.get("top_k", 40))) top_p_val = max(0.0, min(1.0, body.get("top_p", 1.0))) stream = body.get("stream", False) messages = body.get("messages", []) prompt_text = extract_prompt(messages) prompt_tokens = len(tokenizer.encode(prompt_text).ids) if prompt_text else 0 completion_id = f"chatcmpl-{uuid.uuid4()}" created = int(time.time()) if stream: def sse_stream(): initial = { "id": completion_id, "object": "chat.completion.chunk", "created": created, "model": MODEL_ID, "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}], } yield f"data: {json_mod.dumps(initial)}\n\n" token_count = 0 for token_str in generate_streaming( prompt_text, max_tokens=max_tokens, temperature=temperature, top_k=top_k_val, top_p=top_p_val, ): token_count += 1 chunk = { "id": completion_id, "object": "chat.completion.chunk", "created": created, "model": MODEL_ID, "choices": [{"index": 0, "delta": {"content": token_str}, "finish_reason": None}], } yield f"data: {json_mod.dumps(chunk)}\n\n" finish = { "id": completion_id, "object": "chat.completion.chunk", "created": created, "model": MODEL_ID, "choices": [{"index": 0, "delta": {}, "finish_reason": "length" if token_count >= max_tokens else "stop"}], "usage": { "prompt_tokens": prompt_tokens, "completion_tokens": token_count, "total_tokens": prompt_tokens + token_count, }, } yield f"data: {json_mod.dumps(finish)}\n\n" yield "data: [DONE]\n\n" return StreamingResponse(sse_stream(), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}) else: text = generate(prompt_text, max_tokens=max_tokens, temperature=temperature, top_k=top_k_val, top_p=top_p_val) completion_tokens = len(tokenizer.encode(text).ids) return { "id": completion_id, "object": "chat.completion", "created": created, "model": MODEL_ID, "choices": [{ "index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "length", }], "usage": { "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": prompt_tokens + completion_tokens, }, "system_fingerprint": "vocabfusion-pythia160m-fused-v1", } if __name__ == "__main__": print(f"\nVocabFusion-Pythia160m-Fused server starting on 0.0.0.0:{PORT} ...") uvicorn.run(app, host="0.0.0.0", port=PORT)