"""HF Spaces server: static experience plus a bounded, single-model job queue.""" import logging import os import queue import re import threading import time import uuid from contextlib import asynccontextmanager from pathlib import Path from typing import Literal from fastapi import FastAPI, HTTPException from fastapi.staticfiles import StaticFiles from pydantic import BaseModel, Field, field_validator ROOT = Path(__file__).resolve().parent log = logging.getLogger('ouroboros') ORDER = ['similarity', 'function', 'subunit', 'catalytic activity', 'pathway', 'cofactor', 'tissue specificity', 'domain', 'subcellular location', 'PTM', 'miscellaneous', 'induction', 'disruption phenotype', 'activity regulation', 'developmental stage', 'sequence caution', 'caution', 'biophysicochemical properties', 'disease', 'alternative products', 'online information', 'biotechnology', 'polymorphism', 'mass spectrometry', 'allergen', 'toxic dose', 'RNA editing', 'pharmaceutical', 'interaction'] class GenerationRequest(BaseModel): direction: Literal['t2p', 'p2t'] input: str = Field(min_length=1, max_length=8000) temperature: float = Field(default=1, ge=0.1, le=2) greedy: bool = False max_length: int = Field(default=512, ge=32, le=512) @field_validator('input') @classmethod def nonempty(cls, value): if not value.strip(): raise ValueError('Enter a description or a protein sequence.') return value.strip() def normalize_input(request): if request.direction == 't2p': return request.input, None lines = request.input.splitlines() headers = [i for i, s in enumerate(lines) if s.lstrip().startswith('>')] if headers and (len(headers) != 1 or headers[0] != 0): raise HTTPException(422, 'Please submit one protein sequence at a time.') sequence = re.sub(r'\s+', '', ''.join(lines[1:] if headers else lines)).upper() if not sequence or re.search(r'[^ACDEFGHIKLMNPQRSTVWYXBZOU]', sequence): raise HTTPException(422, 'Use amino-acid letters only, optionally with one FASTA header.') if len(sequence) > 511: raise HTTPException(422, 'The original model accepts at most 511 residues. Please shorten the sequence.') normalized = re.sub('[BZOU]', 'X', sequence) return normalized, 'Non-standard residues B/Z/O/U were converted to X.' if normalized != sequence else None class ModelRunner: def __init__(self): self.model = None def load(self): import torch from ouroboros_model import OuroborosModel from huggingface_hub import hf_hub_download torch.set_num_threads(int(os.environ.get('OUROBOROS_CPU_THREADS', '4'))) checkpoint = os.environ.get('CHECKPOINT_PATH') if checkpoint: if not Path(checkpoint).is_file(): raise FileNotFoundError('Configured checkpoint is unavailable') else: checkpoint = hf_hub_download(repo_id=os.environ.get('OUROBOROS_MODEL_REPO', 'REXWind/ouroboros-weights'), filename='dual_model.pt') device = os.environ.get('OUROBOROS_DEVICE', 'cpu') model = OuroborosModel( t5config_path=str(ROOT / 't5config_large.json'), protein_tokenizer_path=os.environ.get('PROTEIN_TOKENIZER_PATH'), text_tokenizer_path=os.environ.get('TEXT_TOKENIZER_PATH'), text_share_param=True, protein_share_param=True, share_fc=True, use_mean_pooling=False, ) model.load_checkpoint(checkpoint) model.to(device).eval() self.model = model def generate(self, request, normalized): import torch if request.direction == 't2p': # Reject silent text truncation before invoking the archived model. tokens = self.model.text_tokenizer(normalized, add_special_tokens=True)['input_ids'] if len(tokens) > 512: raise ValueError('This description exceeds 512 text tokens. Please shorten it.') self.model.max_generation_length = request.max_length args = {'use_greedy': request.greedy, 'temperature': request.temperature, 'top_k': 40, 'top_p': 0.9, 'num_beams': 1 if request.greedy else 4} with torch.inference_mode(): if request.direction == 'p2t': return self.model.protein_to_text([' '.join(normalized)], **args)[0] return self.model.text_to_protein([normalized], num_return_sequences=1, **args)[0] class InferenceEngine: def __init__(self, runner=None, disabled=None): self.runner = runner or ModelRunner() self.disabled = os.environ.get('OUROBOROS_INFERENCE') == 'disabled' if disabled is None else disabled self.state = 'disabled' if self.disabled else 'cold' self.jobs = {} self.lock = threading.RLock() self.pending = queue.Queue(maxsize=5) self.stop = threading.Event() self.thread = None def start(self): if not self.disabled: self.thread = threading.Thread(target=self.work, name='model-worker', daemon=True) self.thread.start() def close(self): self.stop.set() if self.thread: self.thread.join(timeout=1) def submit(self, request): normalized, note = normalize_input(request) if self.disabled: raise HTTPException(503, 'Live inference is unavailable. The Core Idea section contains a recorded demonstration.') with self.lock: now = time.time() self.jobs = {k: v for k, v in self.jobs.items() if now - v['created'] < 1800 or v['status'] in ['queued', 'loading', 'running']} if len(self.jobs) >= 100: oldest = next((k for k, v in self.jobs.items() if v['status'] in ['completed', 'failed']), None) if oldest: del self.jobs[oldest] key = uuid.uuid4().hex job = {'id': key, 'status': 'queued', 'created': now, 'direction': request.direction, 'normalized_input': normalized, 'normalization_note': note, 'model': 'Ouroboros large-v5', 'output': None, 'error': None} self.jobs[key] = job try: self.pending.put_nowait((key, request)) except queue.Full: del self.jobs[key] raise HTTPException(429, 'The model queue is full. Please try again shortly.') return {'id': key, 'status': 'queued'} def read(self, key): with self.lock: if key not in self.jobs: raise HTTPException(404, 'This job does not exist or has expired.') return dict(self.jobs[key]) def work(self): while not self.stop.is_set(): try: key, request = self.pending.get(timeout=0.2) except queue.Empty: continue try: with self.lock: job = self.jobs[key] job['status'] = 'loading' if self.state != 'ready' else 'running' if self.state != 'ready': self.state = 'loading' self.runner.load() self.state = 'ready' with self.lock: job['status'] = 'running' output = self.runner.generate(request, job['normalized_input']) if not output or not output.strip(): raise ValueError('The model returned an empty result. Try another input or sampling.') with self.lock: job.update(status='completed', output=output, finished=time.time()) except Exception as exc: log.exception('Generation job failed') if self.state == 'loading': self.state = 'error' with self.lock: job.update(status='failed', error=str(exc) if isinstance(exc, ValueError) else 'Inference is unavailable. Please retry later.', finished=time.time()) finally: self.pending.task_done() def create_app(engine=None): worker = engine or InferenceEngine() @asynccontextmanager async def lifespan(app): worker.start() yield worker.close() app = FastAPI(title='Ouroboros', lifespan=lifespan) app.state.engine = worker @app.get('/api/health') def health(): return {'status': 'ok', 'model': 'Ouroboros large-v5', 'inference': worker.state, 'queue': worker.pending.qsize()} @app.post('/api/jobs', status_code=202) def submit(request: GenerationRequest): return worker.submit(request) @app.get('/api/jobs/{key}') def get_job(key: str): return worker.read(key) @app.get('/api/prompt-fields') def prompt_fields(): return {'fields': ORDER} if (ROOT / 'dist').is_dir(): app.mount('/', StaticFiles(directory=ROOT / 'dist', html=True), name='frontend') return app app = create_app()