Add BERT metagenome embedding extraction app

Dockerfile

@@ -0,0 +1,27 @@
+FROM python:3.10-slim
+
+ENV PYTHONUNBUFFERED=1
+ENV PYTHONDONTWRITEBYTECODE=1
+ENV TF_CPP_MIN_LOG_LEVEL=2
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user
+ENV PATH=/home/user/.local/bin:$PATH
+
+WORKDIR /home/user/app
+
+COPY --chown=user:user requirements.txt .
+
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r requirements.txt
+
+COPY --chown=user:user . .
+
+EXPOSE 7860
+
+CMD ["python", "app.py"]

README.md

@@ -1,10 +1,36 @@
 ---
-title: Bert Embedding
-emoji: 🔥
-colorFrom: green
+title: BERT Metagenome Embeddings
+emoji: 🧬
+colorFrom: gray
 colorTo: gray
 sdk: docker
 pinned: false
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# bert-embedding
+
+Extract embeddings from DNA sequences using a BERT model pretrained on metagenomic sequences.
+
+## Model
+
+| | |
+|---|---|
+| architecture | BERT, 24 layers, 768 hidden, 12 heads |
+| parameters | ~430M |
+| input | DNA sequence (min 1000 bp) |
+| output | 768-dim embedding |
+| source | [genomenet/bert-metagenome](https://huggingface.co/genomenet/bert-metagenome) |
+
+## Deployment
+
+```bash
+cd /vol/hpcprojects/pmuench/crispr_tool/bert-embedding
+git add -A && git commit -m "update" && git push
+```
+
+## Acknowledgements
+
+- BMBF de.NBI / GenomeNet
+- DFG SPP 2141
+- HZI BIFO

app.py

@@ -0,0 +1,233 @@
+"""
+BERT Metagenome Embeddings - HuggingFace Spaces App
+"""
+
+import os
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+
+import gradio as gr
+import numpy as np
+import tensorflow as tf
+from huggingface_hub import hf_hub_download
+
+# Model config
+MODEL_REPO = "genomenet/bert-metagenome"
+MODEL_FILE = "bert_1k_3.h5"
+WINDOW_SIZE = 1000
+EMBEDDING_LAYER = "layer_transformer_block_21"
+EMBEDDING_DIM = 768
+
+# Singleton model
+_model = None
+_embedding_model = None
+
+def get_model():
+    global _model, _embedding_model
+    if _model is None:
+        print("Downloading model...")
+        model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
+        print(f"Loading model from {model_path}...")
+        _model = tf.keras.models.load_model(model_path, compile=False)
+        _embedding_model = tf.keras.Model(
+            inputs=_model.input,
+            outputs=_model.get_layer(EMBEDDING_LAYER).output
+        )
+        print("Model loaded.")
+    return _embedding_model
+
+def get_gpu_status():
+    gpus = tf.config.list_physical_devices('GPU')
+    if gpus:
+        return f"GPU: {gpus[0].name}"
+    return "CPU only"
+
+# Tokenization
+NUCLEOTIDE_MAP = {
+    'A': [1, 0, 0, 0],
+    'C': [0, 1, 0, 0],
+    'G': [0, 0, 1, 0],
+    'T': [0, 0, 0, 1],
+    'N': [0.25, 0.25, 0.25, 0.25],
+}
+
+def tokenize(sequence):
+    sequence = sequence.upper().replace('U', 'T')
+    tokens = []
+    for char in sequence:
+        if char in NUCLEOTIDE_MAP:
+            tokens.append(NUCLEOTIDE_MAP[char])
+        elif char in 'RYSWKMBDHV':
+            tokens.append(NUCLEOTIDE_MAP['N'])
+    return np.array(tokens, dtype=np.float32)
+
+def validate_sequence(sequence):
+    if not sequence or len(sequence.strip()) == 0:
+        return False, "Sequence is empty"
+    sequence = sequence.upper().replace('U', 'T')
+    valid_chars = set('ACGTNRYSWKMBDHV')
+    invalid = set(sequence) - valid_chars - set(' \n\r\t')
+    if invalid:
+        return False, f"Invalid characters: {invalid}"
+    clean = ''.join(c for c in sequence if c in valid_chars)
+    if len(clean) < WINDOW_SIZE:
+        return False, f"Sequence too short: {len(clean)} < {WINDOW_SIZE} bp"
+    return True, ""
+
+def strip_fasta_header(text):
+    lines = text.strip().split('\n')
+    seq_lines = [l for l in lines if not l.startswith('>')]
+    return ''.join(seq_lines).replace(' ', '').replace('\t', '')
+
+def embed_sequence(sequence, mode="mean", stride=100):
+    """Extract embeddings from sequence."""
+    model = get_model()
+
+    seq_len = len(sequence)
+    embeddings = []
+    positions = []
+
+    for start in range(0, seq_len - WINDOW_SIZE + 1, stride):
+        window = sequence[start:start + WINDOW_SIZE]
+        tokens = tokenize(window)
+        tokens = np.expand_dims(tokens, axis=0)
+
+        emb = model.predict(tokens, verbose=0)
+        embeddings.append(emb[0])
+        positions.append(start)
+
+    embeddings = np.array(embeddings)  # (n_windows, 1000, 768)
+
+    # Pool across sequence positions within each window
+    if mode == "mean":
+        # Mean pool: (n_windows, 768) -> (768,)
+        window_emb = np.mean(embeddings, axis=1)
+        return np.mean(window_emb, axis=0)
+    elif mode == "max":
+        window_emb = np.max(embeddings, axis=1)
+        return np.max(window_emb, axis=0)
+    elif mode == "per-window":
+        # Return per-window mean embeddings
+        return np.mean(embeddings, axis=1)
+    else:
+        return np.mean(np.mean(embeddings, axis=1), axis=0)
+
+# Example sequence
+EXAMPLE_SEQUENCE = """ATGCGATCGATCGATCGATCGTAGCTAGCTAGCTAGCTAGCTGATCGATCGATCGTAGCTAGCTAGCTGATCGATCGATCGATCGATCGTAGCTAGCTAGCTAGCTAGCTGATCGATCGATCGTAGCTAGCTAGCTGATCGATCGATCGATCGATCGTAGCTAGCTAGCTAGCTAGCTGATCGATCGATCGTAGCTAGCTAGCTGATCGATCGATCGATCGATCGTAGCTAGCTAGCTAGCTAGCTGATCGATCGATCGTAGCTAGCTAGCTGATCGATCGATCGATCGATCGTAGCTAGCTAGCTAGCTAGCTGATCGATCGATCGTAGCTAGCTAGCTGATCGATCGATCGATCGATCGTAGCTAGCTAGCTAGCTAGCTGATCGATCGATCGTAGCTAGCTAGCTGATCGATCGATCGATCGATCGTAGCTAGCTAGCTAGCTAGCTGATCGATCGATCGTAGCTAGCTAGCTGATCGATCGATCGATCGATCGTAGCTAGCTAGCTAGCTAGCTGATCGATCGATCGTAGCTAGCTAGCTGATCGATCGATCGATCGATCGTAGCTAGCTAGCTAGCTAGCTGATCGATCGATCGTAGCTAGCTAGCTGATCGATCGATCGATCGATCGTAGCTAGCTAGCTAGCTAGCTGATCGATCGATCGTAGCTAGCTAGCTGATCGATCGATCGATCGATCGTAGCTAGCTAGCTAGCTAGCTGATCGATCGATCGTAGCTAGCTAGCTGATCGATCGATCGATCGATCGTAGCTAGCTAGCTAGCTAGCTGATCGATCGATCGTAGCTAGCTAGCTGATCGATCGATCGATCGATCGTAGCTAGCTAGCTAGCTAGCTGATCGATCGATCGTAGCTAGCTAGCTGATCGATCGATCGATCGATCGTAGCTAGCTAGCTAGCTAGCTGATCGATCGATCGTAGCTAGCTAGCT"""
+
+def process(sequence: str, mode: str, stride: int):
+    """Main processing function."""
+    sequence = strip_fasta_header(sequence.strip())
+
+    is_valid, error = validate_sequence(sequence)
+    if not is_valid:
+        return f"**Error**: {error}", None
+
+    embedding = embed_sequence(sequence, mode=mode, stride=stride)
+
+    if mode == "per-window":
+        # Return as downloadable numpy file
+        import tempfile
+        path = os.path.join(tempfile.gettempdir(), "embeddings.npy")
+        np.save(path, embedding)
+
+        summary = f"""## Embeddings extracted
+
+| | |
+|---|---|
+| sequence length | {len(sequence):,} bp |
+| windows | {embedding.shape[0]} |
+| embedding dim | {embedding.shape[1]} |
+| stride | {stride} bp |
+| shape | {embedding.shape} |
+
+Download the `.npy` file for per-window embeddings.
+"""
+        return summary, path
+    else:
+        # Single vector - show as text
+        emb_str = ", ".join([f"{x:.4f}" for x in embedding[:10]])
+
+        summary = f"""## Embedding extracted
+
+| | |
+|---|---|
+| sequence length | {len(sequence):,} bp |
+| mode | {mode} |
+| embedding dim | {len(embedding)} |
+
+**First 10 dimensions**: [{emb_str}, ...]
+
+Full embedding saved to file.
+"""
+        import tempfile
+        path = os.path.join(tempfile.gettempdir(), "embedding.npy")
+        np.save(path, embedding)
+        return summary, path
+
+# CSS
+CUSTOM_CSS = """
+@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500&display=swap');
+* { font-family: 'Inter', system-ui, sans-serif !important; }
+code, pre, textarea { font-family: 'SF Mono', Consolas, monospace !important; }
+.gradio-container { max-width: 900px !important; background: #fafafa !important; }
+"""
+
+# Build interface
+with gr.Blocks(title="BERT Metagenome Embeddings") as demo:
+    gr.Markdown("""
+# bert-embedding
+
+Extract embeddings from DNA sequences. BERT model (430M params) pretrained on metagenomic sequences.
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            seq_input = gr.Textbox(
+                label="sequence",
+                placeholder="Paste DNA sequence (FASTA or raw)...",
+                lines=8,
+                value=EXAMPLE_SEQUENCE,
+                info="min 1000 bp"
+            )
+            mode_input = gr.Radio(
+                choices=["mean", "max", "per-window"],
+                value="mean",
+                label="pooling",
+                info="mean/max: single 768-dim vector | per-window: (n, 768) matrix"
+            )
+            stride_input = gr.Slider(
+                minimum=50, maximum=500, value=100, step=50,
+                label="stride",
+                info="step size between windows"
+            )
+            btn = gr.Button("extract", variant="primary")
+
+        with gr.Column(scale=1):
+            output = gr.Markdown()
+            download = gr.File(label="download")
+
+    btn.click(
+        process,
+        inputs=[seq_input, mode_input, stride_input],
+        outputs=[output, download]
+    )
+
+    gr.Markdown("""
+---
+**Model**: [genomenet/bert-metagenome](https://huggingface.co/genomenet/bert-metagenome) |
+**Related**: [CRISPR Detection](https://huggingface.co/spaces/genomenet/crispr-array-detection)
+    """)
+
+if __name__ == "__main__":
+    print("Loading model...")
+    model = get_model()
+    print(f"Ready! {get_gpu_status()}")
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        theme=gr.themes.Base(
+            primary_hue=gr.themes.colors.zinc,
+            neutral_hue=gr.themes.colors.zinc,
+        ),
+        css=CUSTOM_CSS
+    )

requirements.txt

@@ -0,0 +1,5 @@
+tensorflow==2.15.1
+keras==2.15.0
+gradio>=4.0.0
+numpy>=1.26.0,<2.0.0
+huggingface_hub>=0.20.0