Uploading initial files

aligned_spaces_desalting.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71e2e2b62eb0b31b864099186f43b1a37b69b1b5579ba0a61d8672abea611d4d
+size 7853137

aligned_spaces_elution.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7632d4d415ab30995513f518799566e5461c2df8845cf37ea7408231cf708b18
+size 7638722

aligned_spaces_lysis.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb0e88bf7d8ad3317534b2d0177310dc082c9f02545c59a155a45c6518641f75
+size 8387736

app.py

@@ -0,0 +1,242 @@
+import streamlit as st
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pickle
+import numpy as np
+import re
+from sklearn.cluster import KMeans
+from transformers import T5Tokenizer, T5EncoderModel
+import os
+from dotenv import load_dotenv
+import google.generativeai as genai
+
+load_dotenv()
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+
+if not GEMINI_API_KEY:
+    st.error("API Key not found! Please check that your .env file exists and is formatted correctly.")
+    st.stop()
+
+genai.configure(api_key=GEMINI_API_KEY)
+
+sys_instruct = """
+You are a Senior Research Biochemist. Your primary role is to design scientifically rigorous, highly cohesive protein purification pipelines based ONLY on provided laboratory data.
+
+Follow these strict guidelines:
+1. Tone & Style: Be concise, professional, and direct. Omit conversational filler and pleasantries. 
+2. Formatting: Use standard Markdown. Create clear headers for each step (e.g., '### Step 1: Lysis'). Bold all specific buffer concentrations, proteins, reagents, and pH values so they are easy to read at the bench.
+3. Biochemical Rationale: For every step, explicitly state *why* specific reagents are used based on standard biochemical principles.
+4. Chemical Guardrails: Never invent or hallucinate protocols or buffers. If you detect chemical incompatibilities in the user's request or retrieved data (e.g., high DTT concentrations applied to standard Ni-NTA columns, or inappropriate detergents for soluble proteins), explicitly flag them with a bold **WARNING**.
+"""
+
+model = genai.GenerativeModel('gemini-3.5-flash', 
+                              system_instruction=sys_instruct)
+
+st.set_page_config(
+    page_title="Bio-CLIP Recommender",
+    layout="wide",
+    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.batch_encode_plus([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")
+    embedder = ProteinEmbedder(device=device)
+    recommender = ProtocolRecommender(device=device)
+    return embedder, recommender
+
+st.title("Bio-CLIP: Protein Purification AI")
+st.markdown("Zero-shot protocol recommendation directly from 1D amino acid sequences.")
+
+with st.spinner("Booting up ProtT5 and Bio-CLIP Models (this takes a moment)..."):
+    try:
+        embedder, recommender = load_ai_engines()
+        models_loaded = True
+    except Exception as e:
+        st.error(f"Failed to load model weights. Ensure `.pth` and `.pkl` files are in the same folder. \n\nError: {e}")
+        models_loaded = False
+
+if models_loaded:
+    with st.sidebar:
+        st.header("About")
+        st.write("This tool uses **ProtT5** to embed amino acids and a Contrastive Learning (**Bio-CLIP**) model to map them to historical purification protocols.")
+        st.markdown("---")
+        st.write("**Top K Results**")
+        top_k = st.slider("Number of recommendations:", 1, 5, 3)
+
+    st.markdown("### Input Sequence")
+    default_gfp = "MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK"
+    sequence_input = st.text_area("Paste Amino Acid Sequence here:", value=default_gfp, height=150)
+
+    user_goal = st.text_input("Optional: What is your specific purification goal? (e.g., 'Prioritize highest yield', 'Are there temperature concerns?')")
+
+    if st.button("Generate Purification Pipeline", type="primary"):
+        sequence_input = sequence_input.strip()
+        
+        if len(sequence_input) < 10:
+            st.warning("Please enter a valid amino acid sequence (at least 10 characters).")
+        else:
+            with st.spinner("1/2: Running ProtT5 Sequence Embedding..."):
+                seq_vector = embedder.embed_raw_sequence(sequence_input)
+            
+            with st.spinner("2/2: Querying Bio-CLIP Latent Space..."):
+                results = recommender.search(seq_vector, top_k=top_k)
+            
+            st.success("Search Complete!")
+            st.markdown("---")
+            
+            col1, col2, col3 = st.columns(3)
+            
+            columns = {'lysis': col1, 'elution': col2, 'desalting': col3}
+            
+            for step in ['lysis', 'elution', 'desalting']:
+                with columns[step]:
+                    st.subheader(f"{step.capitalize()}")
+                    for i, res in enumerate(results[step]):
+                    
+                        conf = max(0.0, min(1.0, (res['confidence'] + 0.1)))
+                        st.markdown(f"**Option {i+1}**")
+                        st.progress(conf, text=f"Confidence: {conf*100:.1f}%")
+                        st.info(res['text'])
+                        st.write("")
+
+            st.markdown("---")
+
+            st.markdown("### 3. AI Synthesis")
+            with st.spinner("3/3: Gemini AI analyzing retrieved protocols..."):
+                try:
+                    retrieved_context = f"""
+                    Lysis Options: {[res['text'] for res in results['lysis']]}
+                    Elution Options: {[res['text'] for res in results['elution']]}
+                    Desalting Options: {[res['text'] for res in results['desalting']]}
+                    """
+                    has_his_tag = "HHHHHH" in sequence_input
+
+                    rag_prompt = f"""
+
+                    User Sequence: {sequence_input}
+                    Contains His-Tag: {has_his_tag}
+                    
+                    Synthesize a cohesive, step-by-step protein purification protocol using ONLY the retrieved options below. 
+                    If the user provided a specific goal, tailor the synthesis to that goal.
+
+                    User's specific goal/question: {user_goal if user_goal else "Provide a standard, recommended cohesive pipeline from these options."}
+
+                    --- RETRIEVED PROTOCOL OPTIONS ---
+                    {retrieved_context}
+                    ----------------------------------
+                    """
+
+                    response = model.generate_content(rag_prompt)
+                    
+                    st.write(response.text)
+                    
+                except Exception as e:
+                    st.error(f"An error occurred while generating the AI summary: {e}")

bio_clip_recommender.py

@@ -0,0 +1,149 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pickle
+import numpy as np
+import re
+from transformers import T5Tokenizer, T5EncoderModel
+
+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)
+        )
+        self.temperature = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    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"):
+        print("Loading ProtT5 Language Model (this takes a moment)...")
+        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()
+        print("ProtT5 Online!")
+
+    def embed_raw_sequence(self, sequence: str):
+        """Converts a raw string of amino acids into the 1024D vector."""
+        seq = re.sub(r"[UZOB]", "X", sequence.upper())
+        
+        seq_spaced = " ".join(list(seq))
+        
+        with torch.no_grad():
+            ids = self.tokenizer.batch_encode_plus([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)
+            embedding = embedding.last_hidden_state
+            
+            seq_len = (attention_mask[0] == 1).sum()
+            protein_emb = embedding[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 = {}
+        
+        print("Loading Bio-CLIP Expert Models...")
+        for step in self.steps:
+            model = BioCLIP().to(self.device)
+            weights_path = f"bioclip_weights_{step}.pth"
+            model.load_state_dict(torch.load(weights_path, map_location=self.device, weights_only=True))
+            
+            model.eval() 
+            self.models[step] = model
+            
+            db_path = f"aligned_spaces_{step}.pkl"
+            with open(db_path, "rb") as f:
+                db = pickle.load(f)
+                self.databases[step] = {
+                    "text_vectors": torch.tensor(db["aligned_text"]).float().to(self.device),
+                    "raw_texts": db["raw_texts"]
+                }
+        print("Ready! All systems online.\n")
+
+    def search(self, protein_sequence_1024d, top_k=3):
+        results = {}
+        
+        if not isinstance(protein_sequence_1024d, torch.Tensor):
+            seq_tensor = torch.tensor(protein_sequence_1024d).float().unsqueeze(0).to(self.device)
+        else:
+            seq_tensor = protein_sequence_1024d.to(self.device)
+            if seq_tensor.dim() == 1:
+                seq_tensor = seq_tensor.unsqueeze(0)
+
+        with torch.no_grad():
+            for step in self.steps:
+                model = self.models[step]
+                db = self.databases[step]
+                
+                z_query = model.seq_projector(seq_tensor)
+                z_query = F.normalize(z_query, p=2, dim=-1)
+                
+                similarities = (z_query @ db["text_vectors"].T).squeeze()
+                
+                top_scores, top_indices = torch.topk(similarities, k=top_k)
+                
+                step_results = []
+                for score, idx in zip(top_scores.cpu().numpy(), top_indices.cpu().numpy()):
+                    step_results.append({
+                        "confidence": score,
+                        "text": db["raw_texts"][idx]
+                    })
+                results[step] = step_results
+                
+        return results
+
+if __name__ == "__main__":
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    try:
+        embedder = ProteinEmbedder(device=device)
+        engine = ProtocolRecommender(device=device)
+        
+        print("\n" + "="*60)
+        raw_amino_acids = "MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK"
+        
+        print(f"User inputted sequence: {raw_amino_acids[:30]}... (Length: {len(raw_amino_acids)})")
+        print("="*60)
+        
+        print("1. Passing sequence through ProtT5...")
+        new_protein_vector = embedder.embed_raw_sequence(raw_amino_acids)
+        
+        print("2. Querying Bio-CLIP databases...\n")
+        recommendations = engine.search(new_protein_vector, top_k=3)
+        
+        print("="*60)
+        print("BIO-CLIP RECOMMENDED PURIFICATION PIPELINE")
+        print("="*60)
+        
+        for step in ['lysis', 'elution', 'desalting']:
+            print(f"\n--- {step.upper()} EXPERT ---")
+            for i, rec in enumerate(recommendations[step]):
+                confidence_pct = max(0, min(100, (rec['confidence'] + 0.1) * 100))
+                
+                print(f"Option {i+1} [Confidence: {confidence_pct:.1f}%]")
+                print(f"Protocol: {rec['text']}\n")
+                
+    except FileNotFoundError as e:
+        print(f"\nERROR: Could not find model files. Make sure you run this in the same folder as your .pth and .pkl files!")
+        print(f"Details: {e}")

bioclip_weights_desalting.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:13a4a32a180a3eb00b9612ae54e182bab32dd32597310d4023781520954c3b88
+size 5779494

bioclip_weights_elution.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:723e1c03c258370db65503a908d51a10d6e3c571acfe6e8eb041c22bee846df5
+size 5779468

bioclip_weights_lysis.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0591cfc98bdaebb86a608dc2efbc29a339a210fa7958f480d850a45eed862e6a
+size 5779378