| 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 |
| from bio_clip_recommender import BioCLIP, ProteinEmbedder, ProtocolRecommender |
|
|
| 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" |
| ) |
|
|
| @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(float(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}") |
