| --- |
| title: Ouroboros |
| emoji: 🧬 |
| colorFrom: blue |
| colorTo: purple |
| sdk: gradio |
| python_version: "3.10" |
| app_file: app.py |
| pinned: false |
| --- |
| |
| # Ouroboros: Bidirectional Protein-Text Generation |
|
|
| Dual T5 model for protein sequence ↔ text description bidirectional generation. |
|
|
| - **Protein → Text**: Given an amino acid sequence, generate a functional description. |
| - **Text → Protein**: Given a text description, generate candidate protein sequences. |
|
|
| --- |
|
|
| ## Architecture |
|
|
| ### Model: `ouroboros_model.py` |
| |
| Two independent T5-large models connected by cross-modal FC projection layers. |
| |
| ``` |
| Protein Modality Text Modality |
| ┌──────────────────────┐ ┌──────────────────────┐ |
| │ p_t5_model (T5-large)│ │ t_t5_model (T5-large)│ |
| │ d_model=1024, 24L │ │ d_model=1024, 24L │ |
| │ vocab=128 (ProtT5) │ │ vocab=32128 (T5-base)│ |
| │ ↑↓ │ │ ↑↓ │ |
| │ p_t5_encoder │ │ t_t5_encoder │ |
| │ ↓ │ │ ↓ │ |
| │ fc_p2t (1024→1024) │←───→│ fc_t2p (1024→1024) │ |
| │ ↓ │ │ ↓ │ |
| │ t_t5_decoder │ │ p_t5_decoder │ |
| └──────────────────────┘ └──────────────────────┘ |
| ``` |
| |
| - **share_fc**: `fc_t2p.weight = transpose(fc_p2t.weight)`, no bias |
| - **text_share_param / protein_share_param**: Encoder-decoder SA weights tied within each T5 |
| - **use_mean_pooling**: False (full sequence cross-attention) |
| - **FP32, ~1.23B params** (after weight sharing) |
| - **Checkpoint**: `dual_model.pt` (2.45 GB) — single full state_dict |
| |
| #### Weight Sharing (`share_weight_t5`) |
| |
| Replicated from `dual_training_t5/utils/share_weight.py:126-159`: |
| - SA q/k/v/o weights: encoder.block[i] ↔ decoder.block[i] |
| - SA layer_norm: encoder.block[i] ↔ decoder.block[i] |
| - FFN DenseReluDense: encoder.block[i] ↔ **decoder.block[0]** (`block[0]` is intentional — mirrors training behavior) |
| - No bias sharing, no FFN layer_norm sharing |
|
|
| #### Tokenizer Handling |
|
|
| | | Protein | Text | |
| |---|---|---| |
| | **Tokenizer** | `Rostlab/prot_t5_xl_uniref50` | `google-t5/t5-base` | |
| | **vocab_size (model)** | 128 (hardcoded) | 32128 (hardcoded) | |
| | **Vocab from tokenizer** | 28 (transformers >=4.40) | 32000 (transformers >=4.40) | |
| |
| **Critical**: Model vocab_sizes are hardcoded, not read from tokenizer. Newer transformers report 28/32000 but checkpoint was trained with 128/32128. Loading from tokenizer causes `RuntimeError: size mismatch`. |
| |
| ### Inference Pipeline |
| |
| ``` |
| p2t: protein_seq → [" ".join(chars)] → protein_tokenizer → p_t5_encoder → fc_p2t → t_t5_model.generate() → text_tokenizer.decode() |
| t2p: text_desc → text_tokenizer → t_t5_encoder → fc_t2p → p_t5_model.generate() → protein_tokenizer.decode() → .replace(" ", "") |
| ``` |
| |
| **Generation params**: temperature=1.0, top_k=40, top_p=0.9, num_beams=4, max_length=512 |
| |
| ### Checkpoint Loading |
| |
| ```python |
| state_dict = torch.load(path, map_location='cpu', weights_only=True) |
| # Strip 'module.' prefix from DDP training |
| new_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()} |
| model.load_state_dict(new_state_dict, strict=False) |
| ``` |
| |
| `strict=False` because shared weights (encoder-decoder tied params, `share_fc` transposed params) may cause missing/unexpected key warnings. |
| |
| --- |
| |
| ## Gradio UI (`app.py`) |
| |
| ### Color Scheme |
| |
| | Element | Protein (red/pink) | Text (blue) | |
| |---------|-------------------|-------------| |
| | Accent | `#E74C3C` | `#2471A3` | |
| | Border | `#E6B0AA` | `#A9CCE3` | |
| | Background | `#FDEDEC` | `#EBF5FB` | |
| |
| ### Tab 1: Protein → Text |
| - Monospace input with red border/focus |
| - Output with blue left-accent border |
| - Advanced: temperature, top_k, top_p, num_beams, greedy toggle |
| |
| ### Tab 2: Text → Protein |
| Three input modes: |
| |
| 1. **Structured Prompt Builder** (default): 29 annotation fields matching UniProt categories |
| - 7 common fields visible (similarity, function, catalytic activity, subunit, cofactor, pathway, subcellular location) |
| - 22 more fields in collapsible accordion |
| - "Build Prompt" button: assembles all fields per ORDER_LIST sequence |
| - "Random Example" button: loads from `100sample.json` |
| |
| 2. **Direct Input**: Collapsible accordion, paste raw text prompt |
| |
| 3. **Assembled Prompt** textbox: always editable, serve as final input to generation |
| |
| ### Prompt Assembly Rule |
| |
| ``` |
| ORDER_LIST traversal → if field has content → append to prompt (no separator) |
| Training format: similarity + function + catalytic activity + path way + ... |
| ``` |
| |
| ### 100sample.json Data |
| |
| JSON array of UniProt samples, each with: |
| - `uniprot_id`, `sequence` (amino acid string) |
| - Annotation fields matching ORDER_LIST keys (each a single-string list) |
| - `prompt`: assembled concatenation of all annotation fields |
| - `func_sim`: short version (similarity + function only) |
| |
| --- |
| |
| ## HF Spaces Deployment |
| |
| ### Current Setup |
| |
| - **Space**: `REXWind/ouroboros` (Gradio SDK, Python 3.10, CPU basic) |
| - **Model Repo**: `REXWind/ouroboros-weights` (stores `dual_model.pt`, ~2.5 GB via Git LFS) |
| - **Auto-download**: `app.py` calls `huggingface_hub.hf_hub_download()` on startup |
| |
| ### Requirements (`requirements.txt`) |
| |
| ``` |
| torch>=2.0.0 |
| transformers>=4.36.0 |
| sentencepiece>=0.1.99 |
| huggingface_hub>=0.20.0 |
| ``` |
| |
| **Do NOT pin gradio** — HF Spaces injects its own gradio version during build. Pinning causes dependency conflicts. |
| |
| ### YAML Front Matter (`README.md`) |
| |
| ```yaml |
| sdk: gradio |
| python_version: "3.10" |
| app_file: app.py |
| ``` |
| |
| - Python 3.10 required (3.13 has pydub/audioop incompatibility) |
| - Gradio SDK version auto-managed by HF |
| |
| --- |
| |
| ## Key Files |
| |
| | File | Purpose | |
| |------|---------| |
| | `app.py` | Gradio web interface (model loading, UI layout, event bindings) | |
| | `ouroboros_model.py` | Self-contained model class (no dependency on training `utils/`) | |
| | `t5config_large.json` | T5-large architecture config | |
| | `100sample.json` | Sample UniProt annotations for Random Example button | |
| | `requirements.txt` | Minimal Python dependencies (gradio managed by HF) | |
|
|
| ### Source Project Paths |
|
|
| All paths on gpfs shared storage: |
| ``` |
| Training code: /mnt/shared-storage-user/dnacoding/protein-text-dual-learn/dual_t5_workspace/dual_training_t5/ |
| Checkpoint: /mnt/shared-storage-user/dnacoding/protein-text-dual-learn/dual_t5_workspace/model_weights/dual_t5_large_newds_v5/best_recovery/dual_model.pt |
| Tokenizer: /mnt/shared-storage-user/dnacoding/protein-text-dual-learn/dual_t5_workspace/hf_weights/prot_t5_xl_uniref50/ |
| Tokenizer: /mnt/shared-storage-user/dnacoding/protein-text-dual-learn/dual_t5_workspace/hf_weights/t5-base/ |
| Demo source: /mnt/shared-storage-user/dnacoding/protein-text-dual-learn/dual_t5_workspace/Ouroboros_hf_demo/ |
| HF Space clone: /mnt/shared-storage-user/dnacoding/protein-text-dual-learn/dual_t5_workspace/hf_space_ouroboros/ |
| ``` |
|
|
| --- |
|
|
| ## Known Issues & Fixes |
|
|
| | Issue | Cause | Fix | |
| |-------|-------|-----| |
| | `size mismatch: [32128] vs [32000]` | New transformers reports vocab=32000 | Hardcode `configs_t.vocab_size = 32128` | |
| | `ImportError: cannot import HfFolder` | gradio 5 + huggingface_hub >=1.0 | Don't pin gradio version in requirements | |
| | `TypeError: unhashable type: 'dict'` | gradio 4.44 + jinja2 >=3.1.5 | Don't pin gradio; use Python 3.10 | |
| | `Cannot install gradio<5.0.0 and gradio==6.15.2` | HF installs gradio 6 | Don't specify gradio in requirements | |
| | `Repository storage limit (Max: 1 GB)` | Space only; model 2.5GB | Put checkpoint in separate Model Repo | |
| | Protein vocab=28 | New transformers prot_t5 vocab | `configs_p.vocab_size = 128` (hardcoded) | |
|
|
