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