File size: 6,433 Bytes
c26bf25 261d7a9 c26bf25 c0618e7 c26bf25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | 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}") |