server.py

"""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()