first try ESM app

app.py

@@ -1,7 +1,245 @@
 import gradio as gr
+import torch
+import numpy as np
+from bio import SeqIO
+import tempfile
+import os
+import json
+from pathlib import Path
+import zipfile
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Load ESM2 model
+print("Loading ESM2 model...")
+import esm
+
+# Load the model and alphabet
+model_name = "esm2_t33_650M_UR50D"
+try:
+    model, alphabet = esm.pretrained.load_model_and_alphabet_local(model_name)
+except:
+    print(f"Loading {model_name} from HuggingFace...")
+    model, alphabet = esm.pretrained.load_model_and_alphabet_hub(model_name)
+
+model = model.eval()
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = model.to(device)
+batch_converter = alphabet.get_batch_converter()
+
+print(f"Model loaded on {device}")
+
+
+def parse_fasta_files(fasta_files):
+    """Parse one or multiple FASTA files and return sequences."""
+    sequences = []
+    file_info = {}
+    
+    for fasta_file in fasta_files:
+        file_name = Path(fasta_file.name).stem
+        file_seqs = []
+        
+        try:
+            for record in SeqIO.parse(fasta_file, "fasta"):
+                sequences.append((record.id, str(record.seq), file_name))
+                file_seqs.append(record.id)
+            file_info[file_name] = file_seqs
+        except Exception as e:
+            raise ValueError(f"Error parsing {fasta_file.name}: {str(e)}")
+    
+    if not sequences:
+        raise ValueError("No sequences found in the provided FASTA files.")
+    
+    return sequences, file_info
+
+
+def generate_embeddings(sequences_batch):
+    """Generate embeddings for a batch of sequences."""
+    # Prepare batch for ESM2
+    batch_labels, batch_strs, batch_tokens = batch_converter(sequences_batch)
+    
+    # Move to device
+    batch_tokens = batch_tokens.to(device)
+    
+    # Generate embeddings
+    with torch.no_grad():
+        results = model(batch_tokens, repr_layers=[33], return_contacts=False)
+    
+    # Extract embeddings (token representations from layer 33)
+    token_embeddings = results["representations"][33]
+    
+    # Get sequence-level embeddings (mean pooling of token embeddings, excluding special tokens)
+    sequence_embeddings = []
+    for i, (label, seq) in enumerate(zip(batch_labels, batch_strs)):
+        # Remove special tokens (first and last)
+        seq_embedding = token_embeddings[i, 1:len(seq) + 1].mean(dim=0)
+        sequence_embeddings.append(seq_embedding.cpu().numpy())
+    
+    return sequence_embeddings
+
+
+def process_embeddings(fasta_files):
+    """Main function to process FASTA files and generate embeddings."""
+    try:
+        # Parse FASTA files
+        sequences, file_info = parse_fasta_files(fasta_files)
+        
+        # Generate embeddings in batches
+        batch_size = 8
+        all_embeddings = {}
+        status_updates = [f"Processing {len(sequences)} sequences from {len(file_info)} file(s)..."]
+        
+        for i in range(0, len(sequences), batch_size):
+            batch = sequences[i:i + batch_size]
+            batch_labels = [(seq_id, seq, file_name) for seq_id, seq, file_name in batch]
+            
+            status_updates.append(f"Generating embeddings for sequences {i + 1}-{min(i + batch_size, len(sequences))}...")
+            
+            # Generate embeddings
+            embeddings = generate_embeddings([(label, seq) for label, seq, _ in batch_labels])
+            
+            # Store embeddings
+            for (seq_id, seq, file_name), embedding in zip(batch_labels, embeddings):
+                key = f"{file_name}_{seq_id}"
+                all_embeddings[key] = {
+                    "sequence_id": seq_id,
+                    "file": file_name,
+                    "sequence_length": len(seq),
+                    "embedding": embedding.tolist()
+                }
+        
+        # Create output files
+        output_files = []
+        temp_dir = tempfile.mkdtemp()
+        
+        # Save embeddings as NPZ (numpy compressed format)
+        npz_path = os.path.join(temp_dir, "embeddings.npz")
+        embeddings_array = {k: np.array(v["embedding"]) for k, v in all_embeddings.items()}
+        np.savez_compressed(npz_path, **embeddings_array)
+        output_files.append(npz_path)
+        status_updates.append(f"Saved compressed embeddings to embeddings.npz")
+        
+        # Save metadata as JSON
+        metadata_path = os.path.join(temp_dir, "metadata.json")
+        metadata = {
+            "num_sequences": len(all_embeddings),
+            "embedding_dim": 1280,  # ESM2-650M has 1280-dimensional embeddings
+            "model": model_name,
+            "sequences": {k: {
+                "sequence_id": v["sequence_id"],
+                "file": v["file"],
+                "sequence_length": v["sequence_length"]
+            } for k, v in all_embeddings.items()}
+        }
+        with open(metadata_path, 'w') as f:
+            json.dump(metadata, f, indent=2)
+        output_files.append(metadata_path)
+        status_updates.append(f"Saved metadata to metadata.json")
+        
+        # Create per-file embedding files
+        for file_name in file_info.keys():
+            file_embeddings = {k: v for k, v in embeddings_array.items() if k.startswith(file_name)}
+            if file_embeddings:
+                file_npz_path = os.path.join(temp_dir, f"embeddings_{file_name}.npz")
+                np.savez_compressed(file_npz_path, **file_embeddings)
+                output_files.append(file_npz_path)
+                status_updates.append(f"Saved {len(file_embeddings)} embeddings for {file_name}")
+        
+        # Create a summary report
+        summary_path = os.path.join(temp_dir, "summary.txt")
+        with open(summary_path, 'w') as f:
+            f.write("ESM2 Protein Sequence Embedding Summary\n")
+            f.write("=" * 50 + "\n\n")
+            f.write(f"Model: {model_name}\n")
+            f.write(f"Device: {device}\n")
+            f.write(f"Embedding Dimension: 1280\n\n")
+            f.write(f"Input Files: {len(file_info)}\n")
+            f.write(f"Total Sequences: {len(all_embeddings)}\n\n")
+            f.write("Sequences per file:\n")
+            for file_name, seq_ids in file_info.items():
+                f.write(f"  - {file_name}: {len(seq_ids)} sequences\n")
+        output_files.append(summary_path)
+        
+        status_message = "\n".join(status_updates)
+        status_message += f"\n\nSuccessfully generated embeddings for {len(all_embeddings)} sequences!"
+        
+        return output_files, status_message
+        
+    except Exception as e:
+        raise gr.Error(f"Error processing sequences: {str(e)}")
+
+
+# Create Gradio interface
+with gr.Blocks(title="ESM2 Protein Embeddings") as demo:
+    gr.Markdown("""
+    # ESM2 Protein Sequence Embeddings
+    
+    Generate embeddings for protein sequences using Meta's ESM2 language model.
+    
+    **Features:**
+    - Process one or multiple FASTA files
+    - Generate high-dimensional embeddings (1280-D) using ESM2-650M
+    - Download embeddings in NumPy format or as JSON metadata
+    - Supports batch processing for efficiency
+    
+    **Instructions:**
+    1. Upload one or more FASTA files containing protein sequences
+    2. Click "Generate Embeddings"
+    3. Download the output files (embeddings.npz, metadata.json, summary.txt)
+    
+    **Output Files:**
+    - `embeddings.npz`: Compressed NumPy file with all embeddings
+    - `metadata.json`: JSON file with sequence IDs and metadata
+    - `summary.txt`: Human-readable summary
+    - `embeddings_[filename].npz`: Per-file embeddings
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            input_files = gr.File(
+                label="Upload FASTA files",
+                file_count="multiple",
+                file_types=[".fasta", ".fa", ".faa"]
+            )
+            submit_btn = gr.Button("Generate Embeddings", variant="primary", size="lg")
+        
+        with gr.Column():
+            status_output = gr.Textbox(
+                label="Processing Status",
+                interactive=False,
+                lines=6
+            )
+    
+    with gr.Row():
+        download_output = gr.File(
+            label="Download Output Files",
+            file_count="multiple"
+        )
+    
+    submit_btn.click(
+        fn=process_embeddings,
+        inputs=[input_files],
+        outputs=[download_output, status_output]
+    )
+    
+    gr.Markdown("""
+    ### How to use the embeddings:
+    
+    ```python
+    import numpy as np
+    import json
+    
+    # Load embeddings
+    embeddings = np.load('embeddings.npz')
+    
+    # Access a specific embedding
+    embedding = embeddings['file_name_sequence_id']
+    
+    # Load metadata
+    with open('metadata.json', 'r') as f:
+        metadata = json.load(f)
+    ```
+    """)
+
+
+if __name__ == "__main__":
+    demo.launch()
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()

requirements.txt

@@ -0,0 +1,5 @@
+gradio==4.26.0
+torch>=2.0.0
+fair-esm>=2.0.0
+biopython>=1.81
+numpy>=1.21.0