hcoops
/

signalseeker

+"""
+SignalSeeker: Protein Signal Peptide Prediction
+"""
+import os
+import json
+import joblib
+import numpy as np
+import torch
+import re
+from transformers import BertModel, BertTokenizer
+from typing import Dict, List, Union, Any
+class SignalSeekerPredictor:
+    """SignalSeeker protein signal peptide predictor"""
+    def __init__(self, model_dir: str):
+        self.model_dir = model_dir
+        self.models = {}
+        self.scaler = None
+        self.protbert_model = None
+        self.protbert_tokenizer = None
+        self.device = None
+        self.config = {}
+        self._load_models()
+    @classmethod
+    def from_pretrained(cls, model_name_or_path: str):
+        """Load model from Hugging Face or local path"""
+        return cls(model_name_or_path)
+    def _load_models(self):
+        """Load all model components"""
+        # Load config
+        config_path = os.path.join(self.model_dir, "config.json")
+        if os.path.exists(config_path):
+            with open(config_path, 'r') as f:
+                self.config = json.load(f)
+        # Load scaler
+        scaler_path = os.path.join(self.model_dir, "scaler.pkl")
+        if os.path.exists(scaler_path):
+            self.scaler = joblib.load(scaler_path)
+        # Load ensemble models
+        models_dir = os.path.join(self.model_dir, "models")
+        if os.path.exists(models_dir):
+            for model_file in os.listdir(models_dir):
+                if model_file.endswith('.pkl'):
+                    model_name = model_file.replace('.pkl', '').replace('_', ' ')
+                    model_path = os.path.join(models_dir, model_file)
+                    self.models[model_name] = joblib.load(model_path)
+        # Load ProtBERT
+        self._load_protbert()
+    def _load_protbert(self):
+        """Load ProtBERT for feature extraction"""
+        try:
+            self.protbert_tokenizer = BertTokenizer.from_pretrained(
+                "Rostlab/prot_bert", do_lower_case=False
+            )
+            self.protbert_model = BertModel.from_pretrained("Rostlab/prot_bert")
+            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+            self.protbert_model = self.protbert_model.to(self.device)
+            self.protbert_model.eval()
+        except Exception as e:
+            raise RuntimeError(f"Failed to load ProtBERT: {e}")
+    def _extract_features(self, sequences: List[str]) -> np.ndarray:
+        """Extract ProtBERT features from sequences"""
+        n_length = self.config.get('n_terminal_length', 50)
+        # Prepare N-terminal sequences
+        n_terminal_seqs = []
+        for seq in sequences:
+            # Clean sequence
+            clean_seq = re.sub(r"[^ACDEFGHIKLMNPQRSTVWXY]", "X", seq)
+            # Extract N-terminal
+            if len(clean_seq) >= n_length:
+                n_terminal_seqs.append(clean_seq[:n_length])
+            else:
+                padded = clean_seq + 'X' * (n_length - len(clean_seq))
+                n_terminal_seqs.append(padded)
+        # Extract embeddings
+        all_embeddings = []
+        batch_size = 8
+        for i in range(0, len(n_terminal_seqs), batch_size):
+            batch = n_terminal_seqs[i:i+batch_size]
+            # Preprocess for ProtBERT
+            prepped = [" ".join(list(seq)) for seq in batch]
+            # Tokenize
+            inputs = self.protbert_tokenizer.batch_encode_plus(
+                prepped,
+                add_special_tokens=True,
+                padding="longest",
+                truncation=True,
+                max_length=512,
+                return_tensors="pt"
+            )
+            input_ids = inputs['input_ids'].to(self.device)
+            attention_mask = inputs['attention_mask'].to(self.device)
+            # Generate embeddings
+            with torch.no_grad():
+                outputs = self.protbert_model(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask
+                )
+            # Extract per-protein embeddings
+            for j, seq in enumerate(batch):
+                seq_length = min(len(seq), 510)
+                if seq_length > 0:
+                    seq_emb = outputs.last_hidden_state[j, 1:seq_length+1]
+                    protein_emb = seq_emb.mean(dim=0).cpu().numpy()
+                else:
+                    protein_emb = outputs.last_hidden_state[j, 0].cpu().numpy()
+                all_embeddings.append(protein_emb)
+        return np.array(all_embeddings)
+    def predict(self, sequence: str) -> Dict[str, Any]:
+        """Predict signal peptide for a single sequence"""
+        # Extract features
+        X = self._extract_features([sequence])
+        # Apply scaling
+        if self.scaler:
+            X = self.scaler.transform(X)
+        # Get predictions from all models
+        predictions = []
+        probabilities = []
+        for model_name, model in self.models.items():
+            try:
+                pred = model.predict(X)[0]
+                prob = model.predict_proba(X)[0, 1]
+                predictions.append(pred)
+                probabilities.append(prob)
+            except Exception:
+                continue
+        # Ensemble decision
+        if probabilities:
+            ensemble_prob = np.mean(probabilities)
+            ensemble_pred = ensemble_prob > 0.5
+            # Confidence assessment
+            if ensemble_prob > 0.8 or ensemble_prob < 0.2:
+                confidence = "High"
+            elif ensemble_prob > 0.7 or ensemble_prob < 0.3:
+                confidence = "Medium"
+            else:
+                confidence = "Low"
+        else:
+            ensemble_prob = 0.0
+            ensemble_pred = False
+            confidence = "Error"
+        return {
+            "has_signal_peptide": bool(ensemble_pred),
+            "probability": float(ensemble_prob),
+            "confidence": confidence,
+            "sequence_length": len(sequence)
+        }
+    def predict_batch(self, sequences: Union[List[str], Dict[str, str]]) -> Dict[str, Dict[str, Any]]:
+        """Predict signal peptides for multiple sequences"""
+        if isinstance(sequences, list):
+            sequences = {f"seq_{i}": seq for i, seq in enumerate(sequences)}
+        results = {}
+        for name, sequence in sequences.items():
+            results[name] = self.predict(sequence)
+        return results
+def pipeline(sequences: Union[str, List[str], Dict[str, str]], **kwargs) -> Union[Dict[str, Any], Dict[str, Dict[str, Any]]]:
+    """Hugging Face pipeline interface"""
+    # This would be automatically set by HF
+    model_dir = kwargs.get('model_dir', '.')
+    predictor = SignalSeekerPredictor(model_dir)
+    if isinstance(sequences, str):
+        return predictor.predict(sequences)
+    else:
+        return predictor.predict_batch(sequences)