updated app.py

app.py

@@ -10,6 +10,7 @@ from transformers import T5Tokenizer, T5EncoderModel
 import os
 from dotenv import load_dotenv
 import google.generativeai as genai
+from bio_clip_recommender import BioCLIP, ProteinEmbedder, ProtocolRecommender
 
 load_dotenv()
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
@@ -39,112 +40,6 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
-class BioCLIP(nn.Module):
-    def __init__(self, seq_dim=1024, text_dim=768, shared_dim=512):
-        super().__init__()
-        self.seq_projector = nn.Sequential(
-            nn.Linear(seq_dim, shared_dim),
-            nn.GELU(),
-            nn.Dropout(0.1),
-            nn.Linear(shared_dim, shared_dim)
-        )
-        self.text_projector = nn.Sequential(
-            nn.Linear(text_dim, shared_dim),
-            nn.GELU(),
-            nn.Dropout(0.1),
-            nn.Linear(shared_dim, shared_dim)
-        )
-
-    def forward(self, seq_emb, text_emb):
-        z_seq = F.normalize(self.seq_projector(seq_emb), p=2, dim=-1)
-        z_text = F.normalize(self.text_projector(text_emb), p=2, dim=-1)
-        return z_seq, z_text
-
-class ProteinEmbedder:
-    def __init__(self, device="cpu"):
-        self.device = device
-        self.tokenizer = T5Tokenizer.from_pretrained("Rostlab/prot_t5_xl_uniref50", do_lower_case=False)
-        self.model = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_uniref50").to(self.device)
-        self.model.eval()
-
-    def embed_raw_sequence(self, sequence: str):
-        seq = re.sub(r"[UZOB]", "X", sequence.upper())
-        seq_spaced = " ".join(list(seq))
-        
-        with torch.no_grad():
-            ids = self.tokenizer([seq_spaced], add_special_tokens=True, padding=True, return_tensors="pt")
-            input_ids = ids['input_ids'].to(self.device)
-            attention_mask = ids['attention_mask'].to(self.device)
-            
-            embedding = self.model(input_ids=input_ids, attention_mask=attention_mask)
-            seq_len = (attention_mask[0] == 1).sum()
-            protein_emb = embedding.last_hidden_state[0, :seq_len-1].mean(dim=0)
-            
-        return protein_emb.cpu().numpy()
-
-class ProtocolRecommender:
-    def __init__(self, device="cpu"):
-        self.device = device
-        self.steps = ['lysis', 'elution', 'desalting']
-        self.models = {}
-        self.databases = {}
-        
-        for step in self.steps:
-            model = BioCLIP().to(self.device)
-            model.load_state_dict(torch.load(f"bioclip_weights_{step}.pth", map_location=self.device, weights_only=True), strict=False)
-            model.eval() 
-            self.models[step] = model
-            
-            with open(f"aligned_spaces_{step}.pkl", "rb") as f:
-                db = pickle.load(f)
-
-                text_vectors_np = np.array(db["aligned_text"])
-                kmeans = KMeans(n_clusters=15, random_state=42, n_init=10)
-                cluster_labels = kmeans.fit_predict(text_vectors_np)
-
-                self.databases[step] = {
-                    "text_vectors": torch.tensor(db["aligned_text"]).float().to(self.device),
-                    "raw_texts": db["raw_texts"],
-                    "cluster_labels": cluster_labels
-                }
-
-    def search(self, protein_sequence_1024d, top_k=3):
-        results = {}
-        seq_tensor = torch.tensor(protein_sequence_1024d).float().unsqueeze(0).to(self.device)
-
-        with torch.no_grad():
-            for step in self.steps:
-                model = self.models[step]
-                db = self.databases[step]
-                
-                z_query = F.normalize(model.seq_projector(seq_tensor), p=2, dim=-1)
-                
-                similarities = (z_query @ db["text_vectors"].T).squeeze().cpu().numpy()
-                
-                sorted_indices = np.argsort(similarities)[::-1]
-                
-                diverse_indices = []
-                seen_clusters = set()
-                
-                for idx in sorted_indices:
-                    cluster_id = db["cluster_labels"][idx]
-                    
-                    if cluster_id not in seen_clusters:
-                        seen_clusters.add(cluster_id)
-                        diverse_indices.append(idx)
-                        
-                        if len(diverse_indices) == top_k:
-                            break
-                            
-                step_results = []
-                for idx in diverse_indices:
-                    step_results.append({
-                        "confidence": float(similarities[idx]),
-                        "text": db["raw_texts"][idx]
-                    })
-                results[step] = step_results
-        return results
-
 @st.cache_resource(show_spinner=False)
 def load_ai_engines():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")