#!/usr/bin/env python3
"""
Tranception Design App - Hugging Face Spaces Version (Zero GPU Fixed)
"""
import os
import sys
# Set up caching to avoid re-downloading models
os.environ['HF_HOME'] = '/tmp/huggingface'
os.environ['TRANSFORMERS_CACHE'] = '/tmp/huggingface/transformers'
os.environ['HF_DATASETS_CACHE'] = '/tmp/huggingface/datasets'
# Ensure proper Hugging Face endpoint
os.environ['HF_ENDPOINT'] = 'https://huggingface.co'
# Disable offline mode to allow downloads
os.environ['TRANSFORMERS_OFFLINE'] = '0'
# Patch for transformers 4.17.0 URL issue in HF Spaces
import urllib.parse
import json
import time
def patch_transformers_url():
"""Fix URL scheme issue in transformers 4.17.0 with comprehensive URL handling"""
try:
import transformers.file_utils
import requests
# Patch multiple functions for comprehensive URL fixing
original_get_from_cache = transformers.file_utils.get_from_cache
def patched_get_from_cache(url, *args, **kwargs):
# Comprehensive URL fixing for various formats
if isinstance(url, str):
original_url = url
# Handle different types of malformed URLs
if url.startswith('/api/resolve-cache/') or url.startswith('/api/'):
# Fix relative API URLs - ensure proper base URL
url = 'https://huggingface.co' + url
print(f"Fixed relative API URL: {original_url} -> {url}")
elif url.startswith('//'):
# Fix protocol-relative URLs
url = 'https:' + url
elif not url.startswith(('http://', 'https://')):
# Handle other relative paths
if url.startswith('/'):
url = 'https://huggingface.co' + url
else:
url = 'https://huggingface.co/' + url
# Additional validation and normalization
try:
parsed = urllib.parse.urlparse(url)
if not parsed.netloc:
# If no netloc found, construct proper URL
url = 'https://huggingface.co' + ('/' + url if not url.startswith('/') else url)
except Exception:
# Fallback for URL parsing errors
if not url.startswith('https://'):
url = 'https://huggingface.co' + ('/' + url if not url.startswith('/') else url)
# Add retry mechanism for network requests
max_retries = 3
for attempt in range(max_retries):
try:
return original_get_from_cache(url, *args, **kwargs)
except Exception as e:
if attempt < max_retries - 1:
print(f"Download attempt {attempt + 1} failed for {url}: {e}. Retrying...")
time.sleep(2 ** attempt) # Exponential backoff
continue
else:
print(f"All download attempts failed for {url}: {e}")
raise
# Also patch cached_path function which might be causing the issue
if hasattr(transformers.file_utils, 'cached_path'):
original_cached_path = transformers.file_utils.cached_path
def patched_cached_path(url_or_filename, *args, **kwargs):
if isinstance(url_or_filename, str):
if url_or_filename.startswith('/api/resolve-cache/') or url_or_filename.startswith('/api/'):
url_or_filename = 'https://huggingface.co' + url_or_filename
print(f"Fixed cached_path URL: {url_or_filename}")
return original_cached_path(url_or_filename, *args, **kwargs)
transformers.file_utils.cached_path = patched_cached_path
# Patch http_get function to handle malformed URLs at the lowest level
if hasattr(transformers.file_utils, 'http_get'):
original_http_get = transformers.file_utils.http_get
def patched_http_get(url, *args, **kwargs):
if isinstance(url, str):
if url.startswith('/api/resolve-cache/') or url.startswith('/api/'):
url = 'https://huggingface.co' + url
print(f"Fixed http_get URL: {url}")
return original_http_get(url, *args, **kwargs)
transformers.file_utils.http_get = patched_http_get
# Patch requests.get at the lowest level to catch any remaining malformed URLs
original_requests_get = requests.get
def patched_requests_get(url, *args, **kwargs):
if isinstance(url, str):
if url.startswith('/api/resolve-cache/') or url.startswith('/api/'):
original_url = url
url = 'https://huggingface.co' + url
print(f"Fixed requests.get URL: {original_url} -> {url}")
elif not url.startswith(('http://', 'https://', 'ftp://')):
if url.startswith('/'):
url = 'https://huggingface.co' + url
print(f"Fixed relative URL in requests.get: {url}")
return original_requests_get(url, *args, **kwargs)
requests.get = patched_requests_get
transformers.file_utils.get_from_cache = patched_get_from_cache
print("Applied comprehensive URL patch for transformers and requests")
except Exception as e:
print(f"Warning: Could not patch transformers URL handling: {e}")
import torch
import transformers
patch_transformers_url()
from transformers import PreTrainedTokenizerFast
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import gradio as gr
from huggingface_hub import hf_hub_download
import shutil
import uuid
import gc
import time
import datetime
import threading
# Simplified Zero GPU handling
try:
import spaces
SPACES_AVAILABLE = True
print("Zero GPU support detected")
except ImportError:
SPACES_AVAILABLE = False
print("Running without Zero GPU support")
except Exception as e:
# Catch any other initialization errors
SPACES_AVAILABLE = False
print(f"Zero GPU initialization warning: {e}")
print("Running without Zero GPU support")
# Runtime mode tracking
RUNTIME_MODE = "GPU" if SPACES_AVAILABLE else "CPU"
# Keep-alive state
last_activity = datetime.datetime.now()
activity_lock = threading.Lock()
def update_activity():
"""Update last activity timestamp"""
global last_activity
with activity_lock:
last_activity = datetime.datetime.now()
# Add current directory to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
# Check if we need to download and extract the tranception module
if not os.path.exists("tranception"):
print("Downloading Tranception repository...")
try:
# Clone the repository structure
result = os.system("git clone https://github.com/OATML-Markslab/Tranception.git temp_tranception")
if result != 0:
raise Exception("Failed to clone Tranception repository")
# Move the tranception module to current directory
shutil.move("temp_tranception/tranception", "tranception")
# Clean up
shutil.rmtree("temp_tranception")
except Exception as e:
print(f"Error setting up Tranception: {e}")
if os.path.exists("temp_tranception"):
shutil.rmtree("temp_tranception")
raise
import tranception
from tranception import config, model_pytorch
# Model loading configuration
MODEL_CACHE = {}
def validate_cache_file(file_path, min_size=1000):
"""Validate cache file integrity and content"""
if not os.path.exists(file_path):
return False, "File does not exist"
# Check file size
try:
file_size = os.path.getsize(file_path)
if file_size < min_size:
return False, f"File too small ({file_size} bytes < {min_size})"
except Exception as e:
return False, f"Cannot get file size: {e}"
# Check if it's supposed to be a JSON file (config files)
if file_path.endswith('.json') or 'config' in file_path.lower():
try:
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read().strip()
if not content:
return False, "Empty JSON file"
json.loads(content) # Validate JSON syntax
return True, "Valid JSON file"
except json.JSONDecodeError:
return False, "Invalid JSON content"
except Exception as e:
return False, f"Cannot read JSON file: {e}"
return True, "File appears valid"
def clean_corrupted_cache_files(cache_dir):
"""Clean corrupted or invalid cache files"""
if not os.path.exists(cache_dir):
return
cleaned_count = 0
for file in os.listdir(cache_dir):
filepath = os.path.join(cache_dir, file)
if os.path.isfile(filepath):
valid, reason = validate_cache_file(filepath)
if not valid:
try:
os.remove(filepath)
print(f"Removed corrupted cache file: {file} ({reason})")
cleaned_count += 1
except Exception as e:
print(f"Could not remove {file}: {e}")
if cleaned_count > 0:
print(f"Cleaned {cleaned_count} corrupted cache files")
def get_model_path(model_name):
"""Get model path - always use HF Hub for Zero GPU spaces"""
# In HF Spaces, models are cached automatically by the transformers library
# Always return the HF Hub path to leverage this caching
return f"PascalNotin/{model_name}"
def load_model_direct(model_type):
"""Direct model loading with manual config handling"""
import json
import tempfile
from transformers import AutoConfig
print(f"Attempting direct load of {model_type} model...")
# Create a proper config manually based on model type
config_data = {
"architectures": ["TranceptionLMHeadModel"],
"model_type": "tranception",
"_name_or_path": f"Tranception_{model_type}",
"activation_function": "squared_relu",
"attention_mode": "tranception",
"attn_pdrop": 0.1,
"embd_pdrop": 0.1,
"initializer_range": 0.02,
"layer_norm_epsilon": 1e-5,
"n_embd": 768 if model_type == "Small" else (1024 if model_type == "Medium" else 1280),
"n_head": 12 if model_type == "Small" else (16 if model_type == "Medium" else 20),
"n_inner": None,
"n_layer": 12 if model_type == "Small" else (24 if model_type == "Medium" else 30),
"n_positions": 2048,
"resid_pdrop": 0.1,
"summary_activation": None,
"summary_first_dropout": 0.1,
"summary_proj_to_labels": True,
"summary_type": "cls_index",
"summary_use_proj": True,
"vocab_size": 50257,
"pad_token_id": 50256,
"bos_token_id": 50256,
"eos_token_id": 50256
}
# Save config to temp file
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
json.dump(config_data, f)
config_path = f.name
try:
# Load config from temp file
try:
config = AutoConfig.from_pretrained(config_path)
except Exception:
# Try without trust_remote_code
config = AutoConfig.from_pretrained(config_path)
# Load model with manual config
try:
model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
f"PascalNotin/Tranception_{model_type}",
config=config,
ignore_mismatched_sizes=True
)
except TypeError:
# Fallback without newer parameters
model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
f"PascalNotin/Tranception_{model_type}",
config=config
)
os.unlink(config_path) # Clean up temp file
return model
except Exception as e:
print(f"Direct load failed: {e}")
if os.path.exists(config_path):
os.unlink(config_path)
raise
def load_model_cached(model_type):
"""Load model with caching to avoid re-downloading"""
global MODEL_CACHE
# Check if model is already in cache
if model_type in MODEL_CACHE:
print(f"Using cached {model_type} model")
return MODEL_CACHE[model_type]
print(f"Loading {model_type} model...")
model_name = f"Tranception_{model_type}"
model_path = get_model_path(model_name)
try:
# Enhanced cache cleaning with validation
import shutil
cache_dir = "/tmp/huggingface/transformers"
os.makedirs(cache_dir, exist_ok=True)
# Clean corrupted cache files using the new validation system
print("Validating and cleaning cache files...")
clean_corrupted_cache_files(cache_dir)
# Enhanced environment setup for robust model loading
os.environ["HF_ENDPOINT"] = "https://huggingface.co"
os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
os.environ["HF_HUB_DISABLE_EXPERIMENTAL_WARNING"] = "1"
# Try loading without trust_remote_code first (compatibility issue)
try:
model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
model_path,
cache_dir=cache_dir,
force_download=True,
resume_download=False
)
except Exception as e1:
print(f"Loading without trust_remote_code failed: {e1}")
# Fallback: try with trust_remote_code for older transformers versions
try:
model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
model_path,
cache_dir=cache_dir,
force_download=True,
trust_remote_code=True,
resume_download=False
)
except Exception as e2:
print(f"Loading with trust_remote_code also failed: {e2}")
raise e1 # Raise the original exception
MODEL_CACHE[model_type] = model
print(f"{model_type} model loaded and cached")
return model
except Exception as e:
print(f"Error loading {model_type} model: {e}")
print(f"Error type: {type(e).__name__}")
if hasattr(e, '__cause__') and e.__cause__:
print(f"Root cause: {e.__cause__}")
print(f"Model path used: {model_path}")
print(f"Cache directory: {cache_dir}")
print(f"Attempting alternative loading method...")
# Try alternative loading approach with full URL
try:
# Use full URL to bypass any path resolution issues
full_url = f"https://huggingface.co/PascalNotin/Tranception_{model_type}"
try:
model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
full_url,
cache_dir=cache_dir
)
except TypeError:
# Try without trust_remote_code if it's not supported
model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
full_url,
cache_dir=cache_dir
)
MODEL_CACHE[model_type] = model
print(f"{model_type} model loaded successfully with full URL")
return model
except Exception as e2:
print(f"Alternative loading also failed: {e2}")
# Final attempt: use AutoModel with manual config
try:
import json
import requests
from transformers import AutoConfig, AutoModel
print(f"Attempting to load with AutoModel...")
# Clear cache and try with AutoModel which handles config better
cache_dir_auto = "/tmp/huggingface/auto"
os.makedirs(cache_dir_auto, exist_ok=True)
# Try direct loading with manual config
model = load_model_direct(model_type)
MODEL_CACHE[model_type] = model
print(f"{model_type} model loaded successfully with AutoConfig")
return model
except Exception as e3:
print(f"AutoModel loading also failed: {e3}")
# Fallback to Medium if requested model fails
if model_type == "Large":
print("Falling back to Medium model...")
return load_model_cached("Medium")
elif model_type == "Medium":
print("Medium model failed, trying Small model...")
# Try Small model as last resort
try:
try:
model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
"PascalNotin/Tranception_Small",
force_download=True,
cache_dir="/tmp/huggingface/small"
)
except TypeError:
model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(
"PascalNotin/Tranception_Small",
force_download=True,
cache_dir="/tmp/huggingface/small"
)
MODEL_CACHE["Small"] = model
print("Small model loaded as fallback")
return model
except Exception as e_small:
print(f"Small model also failed: {e_small}")
raise RuntimeError("Failed to load any Tranception model")
else:
raise RuntimeError(f"Failed to load {model_type} model")
AA_vocab = "ACDEFGHIKLMNPQRSTVWY"
tokenizer = PreTrainedTokenizerFast(tokenizer_file="./tranception/utils/tokenizers/Basic_tokenizer",
unk_token="[UNK]",
sep_token="[SEP]",
pad_token="[PAD]",
cls_token="[CLS]",
mask_token="[MASK]"
)
def create_all_single_mutants(sequence,AA_vocab=AA_vocab,mutation_range_start=None,mutation_range_end=None):
all_single_mutants={}
sequence_list=list(sequence)
if mutation_range_start is None: mutation_range_start=1
if mutation_range_end is None: mutation_range_end=len(sequence)
for position,current_AA in enumerate(sequence[mutation_range_start-1:mutation_range_end]):
for mutated_AA in AA_vocab:
if current_AA!=mutated_AA:
mutated_sequence = sequence_list.copy()
mutated_sequence[mutation_range_start + position - 1] = mutated_AA
all_single_mutants[current_AA+str(mutation_range_start+position)+mutated_AA]="".join(mutated_sequence)
all_single_mutants = pd.DataFrame.from_dict(all_single_mutants,columns=['mutated_sequence'],orient='index')
all_single_mutants.reset_index(inplace=True)
all_single_mutants.columns = ['mutant','mutated_sequence']
return all_single_mutants
def create_scoring_matrix_visual(scores,sequence,image_index=0,mutation_range_start=None,mutation_range_end=None,AA_vocab=AA_vocab,annotate=True,fontsize=20,unique_id=None):
if unique_id is None:
unique_id = str(uuid.uuid4())
filtered_scores=scores.copy()
filtered_scores=filtered_scores[filtered_scores.position.isin(range(mutation_range_start,mutation_range_end+1))]
piv=filtered_scores.pivot(index='position',columns='target_AA',values='avg_score').round(4)
# Calculate mutation range length
mutation_range_len = mutation_range_end - mutation_range_start + 1
# Save CSV file
csv_path = 'fitness_scoring_substitution_matrix_{}_{}.csv'.format(unique_id, image_index)
# Create a more detailed CSV with mutation info
csv_data = []
for position in range(mutation_range_start,mutation_range_end+1):
for target_AA in list(AA_vocab):
mutant = sequence[position-1]+str(position)+target_AA
if mutant in set(filtered_scores.mutant):
score_value = filtered_scores.loc[filtered_scores.mutant==mutant,'avg_score']
if isinstance(score_value, pd.Series):
score = float(score_value.iloc[0])
else:
score = float(score_value)
else:
score = 0.0
csv_data.append({
'position': position,
'original_AA': sequence[position-1],
'target_AA': target_AA,
'mutation': mutant,
'fitness_score': score
})
csv_df = pd.DataFrame(csv_data)
csv_df.to_csv(csv_path, index=False)
# Continue with visualization
# Use large fixed width for clarity, height scales with positions (as in reference)
fig, ax = plt.subplots(figsize=(50, mutation_range_len))
scores_dict = {}
valid_mutant_set=set(filtered_scores.mutant)
ax.tick_params(bottom=True, top=True, left=True, right=True)
ax.tick_params(labelbottom=True, labeltop=True, labelleft=True, labelright=True)
if annotate:
for position in range(mutation_range_start,mutation_range_end+1):
for target_AA in list(AA_vocab):
mutant = sequence[position-1]+str(position)+target_AA
if mutant in valid_mutant_set:
score_value = filtered_scores.loc[filtered_scores.mutant==mutant,'avg_score']
if isinstance(score_value, pd.Series):
scores_dict[mutant] = float(score_value.iloc[0])
else:
scores_dict[mutant] = float(score_value)
else:
scores_dict[mutant]=0.0
# Format labels as in reference - always show mutation and score with 4 decimal places
labels = (np.asarray(["{} \n {:.4f}".format(symb,value) for symb, value in scores_dict.items() ])).reshape(mutation_range_len,len(AA_vocab))
heat = sns.heatmap(piv,annot=labels,fmt="",cmap='RdYlGn',linewidths=0.30,ax=ax,vmin=np.percentile(scores.avg_score,2),vmax=np.percentile(scores.avg_score,98),\
cbar_kws={'label': 'Log likelihood ratio (mutant / starting sequence)'},annot_kws={"size": fontsize})
else:
heat = sns.heatmap(piv,cmap='RdYlGn',linewidths=0.30,ax=ax,vmin=np.percentile(scores.avg_score,2),vmax=np.percentile(scores.avg_score,98),\
cbar_kws={'label': 'Log likelihood ratio (mutant / starting sequence)'},annot_kws={"size": fontsize})
# Use label sizes from reference
heat.figure.axes[-1].yaxis.label.set_size(fontsize=int(fontsize*1.5))
heat.set_title("Higher predicted scores (green) imply higher protein fitness",fontsize=fontsize*2, pad=40)
heat.set_ylabel("Sequence position", fontsize = fontsize*2)
heat.set_xlabel("Amino Acid mutation", fontsize = fontsize*2)
# Set y-axis labels (positions)
yticklabels = [str(pos)+' ('+sequence[pos-1]+')' for pos in range(mutation_range_start,mutation_range_end+1)]
heat.set_yticklabels(yticklabels, fontsize=fontsize, rotation=0)
# Set x-axis labels (amino acids) - ensuring correct number
heat.set_xticklabels(list(AA_vocab), fontsize=fontsize)
try:
plt.tight_layout()
image_path = 'fitness_scoring_substitution_matrix_{}_{}.png'.format(unique_id, image_index)
plt.savefig(image_path, dpi=100)
return image_path, csv_path
finally:
plt.close('all') # Ensure all figures are closed
plt.clf() # Clear the current figure
plt.cla() # Clear the current axes
def suggest_mutations(scores):
intro_message = "The following mutations may be sensible options to improve fitness: \n\n"
#Best mutants
top_mutants=list(scores.sort_values(by=['avg_score'],ascending=False).head(5).mutant)
top_mutants_fitness=list(scores.sort_values(by=['avg_score'],ascending=False).head(5).avg_score)
top_mutants_recos = [top_mutant+" ("+str(round(top_mutant_fitness,4))+")" for (top_mutant,top_mutant_fitness) in zip(top_mutants,top_mutants_fitness)]
mutant_recos = "The single mutants with highest predicted fitness are (positive scores indicate fitness increase Vs starting sequence, negative scores indicate fitness decrease):\n {} \n\n".format(", ".join(top_mutants_recos))
#Best positions
positive_scores = scores[scores.avg_score > 0]
if len(positive_scores) > 0:
# Only select numeric columns for groupby mean
positive_scores_position_avg = positive_scores.groupby(['position'])['avg_score'].mean().reset_index()
top_positions=list(positive_scores_position_avg.sort_values(by=['avg_score'],ascending=False).head(5)['position'].astype(str))
position_recos = "The positions with the highest average fitness increase are (only positions with at least one fitness increase are considered):\n {}".format(", ".join(top_positions))
else:
position_recos = "No positions with positive fitness effects found."
return intro_message+mutant_recos+position_recos
def check_valid_mutant(sequence,mutant,AA_vocab=AA_vocab):
valid = True
try:
from_AA, position, to_AA = mutant[0], int(mutant[1:-1]), mutant[-1]
except:
valid = False
if valid and position > 0 and position <= len(sequence):
if sequence[position-1]!=from_AA: valid=False
else:
valid = False
if to_AA not in AA_vocab: valid=False
return valid
def cleanup_old_files(max_age_minutes=30):
"""Clean up old inference files"""
import glob
current_time = time.time()
patterns = ["fitness_scoring_substitution_matrix_*.png",
"fitness_scoring_substitution_matrix_*.csv",
"all_mutations_fitness_scores_*.csv"]
cleaned_count = 0
for pattern in patterns:
for file_path in glob.glob(pattern):
try:
file_age = current_time - os.path.getmtime(file_path)
if file_age > max_age_minutes * 60:
os.remove(file_path)
cleaned_count += 1
except Exception as e:
# Log error but continue cleaning other files
print(f"Warning: Could not remove {file_path}: {e}")
if cleaned_count > 0:
print(f"Cleaned up {cleaned_count} old files")
def get_mutated_protein(sequence,mutant):
if not check_valid_mutant(sequence,mutant):
return "The mutant is not valid"
mutated_sequence = list(sequence)
mutated_sequence[int(mutant[1:-1])-1]=mutant[-1]
return ''.join(mutated_sequence)
def score_and_create_matrix_all_singles_impl(sequence,mutation_range_start=None,mutation_range_end=None,model_type="Large",scoring_mirror=False,batch_size_inference=20,max_number_positions_per_heatmap=50,num_workers=0,AA_vocab=AA_vocab):
# Update activity
update_activity()
# Clean up old files periodically
cleanup_old_files()
# Generate unique ID for this request
unique_id = str(uuid.uuid4())
if mutation_range_start is None: mutation_range_start=1
if mutation_range_end is None: mutation_range_end=len(sequence)
# Clean sequence
sequence = sequence.strip().upper()
# Validate
assert len(sequence) > 0, "no sequence entered"
assert mutation_range_start <= mutation_range_end, "mutation range is invalid"
assert mutation_range_end <= len(sequence), f"End position ({mutation_range_end}) exceeds sequence length ({len(sequence)})"
# Load model with caching
model = load_model_cached(model_type)
# Move model to appropriate device INSIDE the GPU decorated function
# This is crucial for Zero GPU - the model must be moved to GPU inside the decorated function
# Device selection - Zero GPU will provide CUDA when decorated with @spaces.GPU
print(f"GPU Available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
device = torch.device("cuda")
model = model.to(device)
gpu_name = torch.cuda.get_device_name(0)
gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
print(f"Inference will take place on {gpu_name}")
print(f"GPU Memory: {gpu_memory:.2f} GB")
# Increase batch size for GPU inference
batch_size_inference = min(batch_size_inference, 50)
else:
device = torch.device("cpu")
model = model.to(device)
print("Inference will take place on CPU")
# Reduce batch size for CPU inference
batch_size_inference = min(batch_size_inference, 10)
try:
model.eval()
model.config.tokenizer = tokenizer
all_single_mutants = create_all_single_mutants(sequence,AA_vocab,mutation_range_start,mutation_range_end)
with torch.no_grad():
scores = model.score_mutants(DMS_data=all_single_mutants,
target_seq=sequence,
scoring_mirror=scoring_mirror,
batch_size_inference=batch_size_inference,
num_workers=num_workers,
indel_mode=False
)
scores = pd.merge(scores,all_single_mutants,on="mutated_sequence",how="left")
scores["position"]=scores["mutant"].map(lambda x: int(x[1:-1]))
scores["target_AA"] = scores["mutant"].map(lambda x: x[-1])
score_heatmaps = []
csv_files = []
mutation_range = mutation_range_end - mutation_range_start + 1
number_heatmaps = int((mutation_range - 1) / max_number_positions_per_heatmap) + 1
image_index = 0
window_start = mutation_range_start
window_end = min(mutation_range_end,mutation_range_start+max_number_positions_per_heatmap-1)
for image_index in range(number_heatmaps):
image_path, csv_path = create_scoring_matrix_visual(scores,sequence,image_index,window_start,window_end,AA_vocab,unique_id=unique_id)
score_heatmaps.append(image_path)
csv_files.append(csv_path)
window_start += max_number_positions_per_heatmap
window_end = min(mutation_range_end,window_start+max_number_positions_per_heatmap-1)
# Also save a comprehensive CSV with all mutations
comprehensive_csv_path = 'all_mutations_fitness_scores_{}.csv'.format(unique_id)
scores_export = scores[['mutant', 'position', 'target_AA', 'avg_score', 'mutated_sequence']].copy()
scores_export['original_AA'] = scores_export['mutant'].str[0]
scores_export = scores_export.rename(columns={'avg_score': 'fitness_score'})
scores_export = scores_export[['position', 'original_AA', 'target_AA', 'mutant', 'fitness_score', 'mutated_sequence']]
scores_export.to_csv(comprehensive_csv_path, index=False)
csv_files.append(comprehensive_csv_path)
return score_heatmaps, suggest_mutations(scores), csv_files
finally:
# Clean up GPU memory but keep model in cache
# Move model back to CPU to free GPU memory
if 'model' in locals():
model.cpu()
if torch.cuda.is_available():
torch.cuda.empty_cache()
gc.collect()
# Apply Zero GPU decorator if available
if SPACES_AVAILABLE:
try:
score_and_create_matrix_all_singles = spaces.GPU(duration=420)(score_and_create_matrix_all_singles_impl)
except Exception as e:
print(f"Warning: Could not apply Zero GPU decorator: {e}")
print("Falling back to CPU mode")
score_and_create_matrix_all_singles = score_and_create_matrix_all_singles_impl
else:
score_and_create_matrix_all_singles = score_and_create_matrix_all_singles_impl
def extract_sequence(protein_id, taxon, sequence):
return sequence
def clear_inputs(protein_sequence_input,mutation_range_start,mutation_range_end):
protein_sequence_input = ""
mutation_range_start = None
mutation_range_end = None
return protein_sequence_input,mutation_range_start,mutation_range_end
# Create Gradio app
tranception_design = gr.Blocks()
with tranception_design:
gr.Markdown("# In silico directed evolution for protein redesign with Tranception")
gr.Markdown("## 🧬 BASIS-China iGEM Team 2025 - Protein Engineering Platform")
gr.Markdown("### Welcome to BASIS-China's implementation of Tranception on Hugging Face Spaces!")
gr.Markdown("We are the BASIS-China iGEM team, and we're excited to present our deployment of the Tranception model for protein fitness prediction. This tool enables in silico directed evolution to iteratively improve protein fitness through single amino acid substitutions. At each step, Tranception computes log likelihood ratios for all possible mutations compared to the starting sequence, generating fitness heatmaps and recommendations to guide protein engineering.")
gr.Markdown("**Technical Details**: This deployment leverages Hugging Face's Zero GPU infrastructure, which dynamically allocates H200 GPU resources when available. This allows for efficient inference while managing computational resources effectively.")
# Hidden keep-alive component
with gr.Row(visible=False):
keep_alive_component = gr.Number(value=0, visible=False)
def keep_alive_update():
update_activity()
return time.time()
# Update every 2 minutes to keep websocket alive
keep_alive_timer = gr.Timer(value=120)
keep_alive_timer.tick(keep_alive_update, outputs=[keep_alive_component])
# Status indicator
with gr.Row():
with gr.Column(scale=1):
def get_gpu_status():
global RUNTIME_MODE
with activity_lock:
time_since = (datetime.datetime.now() - last_activity).total_seconds()
if RUNTIME_MODE == "GPU":
status = "🔥 Zero GPU"
else:
status = "💻 CPU Mode (GPU initialization failed)"
return f"{status} | Last activity: {int(time_since)}s ago"
gpu_status = gr.Textbox(
label="Compute Status",
value=get_gpu_status,
every=5, # Update every 5 seconds
interactive=False,
elem_id="gpu_status"
)
with gr.Tabs():
with gr.TabItem("Input"):
with gr.Row():
protein_sequence_input = gr.Textbox(lines=1,
label="Protein sequence",
placeholder = "Input the sequence of amino acids representing the starting protein of interest or select one from the list of examples below. You may enter the full sequence or just a subdomain (providing full context typically leads to better results, but is slower at inference)"
)
with gr.Row():
mutation_range_start = gr.Number(label="Start of mutation window (first position indexed at 1)", value=1, precision=0)
mutation_range_end = gr.Number(label="End of mutation window (leave empty for full lenth)", value=10, precision=0)
with gr.TabItem("Parameters"):
with gr.Row():
model_size_selection = gr.Radio(label="Tranception model size (larger models are more accurate but are slower at inference)",
choices=["Small","Medium","Large"],
value="Small")
with gr.Row():
scoring_mirror = gr.Checkbox(label="Score protein from both directions (leads to more robust fitness predictions, but doubles inference time)")
with gr.Row():
batch_size_inference = gr.Number(label="Model batch size at inference time (reduce for CPU)",value = 10, precision=0)
with gr.Row():
gr.Markdown("Note: the current version does not leverage retrieval of homologs at inference time to increase fitness prediction performance.")
with gr.Row():
clear_button = gr.Button(value="Clear",variant="secondary")
run_button = gr.Button(value="Predict fitness",variant="primary")
protein_ID = gr.Textbox(label="Uniprot ID", visible=False)
taxon = gr.Textbox(label="Taxon", visible=False)
examples = gr.Examples(
inputs=[protein_ID, taxon, protein_sequence_input],
outputs=[protein_sequence_input],
fn=extract_sequence,
examples=[
['ADRB2_HUMAN' ,'Human', 'MGQPGNGSAFLLAPNGSHAPDHDVTQERDEVWVVGMGIVMSLIVLAIVFGNVLVITAIAKFERLQTVTNYFITSLACADLVMGLAVVPFGAAHILMKMWTFGNFWCEFWTSIDVLCVTASIETLCVIAVDRYFAITSPFKYQSLLTKNKARVIILMVWIVSGLTSFLPIQMHWYRATHQEAINCYANETCCDFFTNQAYAIASSIVSFYVPLVIMVFVYSRVFQEAKRQLQKIDKSEGRFHVQNLSQVEQDGRTGHGLRRSSKFCLKEHKALKTLGIIMGTFTLCWLPFFIVNIVHVIQDNLIRKEVYILLNWIGYVNSGFNPLIYCRSPDFRIAFQELLCLRRSSLKAYGNGYSSNGNTGEQSGYHVEQEKENKLLCEDLPGTEDFVGHQGTVPSDNIDSQGRNCSTNDSLL'],
['IF1_ECOLI' ,'Prokaryote', 'MAKEDNIEMQGTVLETLPNTMFRVELENGHVVTAHISGKMRKNYIRILTGDKVTVELTPYDLSKGRIVFRSR'],
['P53_HUMAN' ,'Human', 'MEEPQSDPSVEPPLSQETFSDLWKLLPENNVLSPLPSQAMDDLMLSPDDIEQWFTEDPGPDEAPRMPEAAPRVAPAPAAPTPAAPAPAPSWPLSSSVPSQKTYQGSYGFRLGFLHSGTAKSVTCTYSPALNKMFCQLAKTCPVQLWVDSTPPPGTRVRAMAIYKQSQHMTEVVRRCPHHERCSDSDGLAPPQHLIRVEGNLRVEYLDDRNTFRHSVVVPYEPPEVGSDCTTIHYNYMCNSSCMGGMNRRPILTIITLEDSSGNLLGRNSFEVRVCACPGRDRRTEEENLRKKGEPHHELPPGSTKRALPNNTSSSPQPKKKPLDGEYFTLQIRGRERFEMFRELNEALELKDAQAGKEPGGSRAHSSHLKSKKGQSTSRHKKLMFKTEGPDSD'],
['BLAT_ECOLX' ,'Prokaryote', 'MSIQHFRVALIPFFAAFCLPVFAHPETLVKVKDAEDQLGARVGYIELDLNSGKILESFRPEERFPMMSTFKVLLCGAVLSRVDAGQEQLGRRIHYSQNDLVEYSPVTEKHLTDGMTVRELCSAAITMSDNTAANLLLTTIGGPKELTAFLHNMGDHVTRLDRWEPELNEAIPNDERDTTMPAAMATTLRKLLTGELLTLASRQQLIDWMEADKVAGPLLRSALPAGWFIADKSGAGERGSRGIIAALGPDGKPSRIVVIYTTGSQATMDERNRQIAEIGASLIKHW'],
['BRCA1_HUMAN' ,'Human', 'MDLSALRVEEVQNVINAMQKILECPICLELIKEPVSTKCDHIFCKFCMLKLLNQKKGPSQCPLCKNDITKRSLQESTRFSQLVEELLKIICAFQLDTGLEYANSYNFAKKENNSPEHLKDEVSIIQSMGYRNRAKRLLQSEPENPSLQETSLSVQLSNLGTVRTLRTKQRIQPQKTSVYIELGSDSSEDTVNKATYCSVGDQELLQITPQGTRDEISLDSAKKAACEFSETDVTNTEHHQPSNNDLNTTEKRAAERHPEKYQGSSVSNLHVEPCGTNTHASSLQHENSSLLLTKDRMNVEKAEFCNKSKQPGLARSQHNRWAGSKETCNDRRTPSTEKKVDLNADPLCERKEWNKQKLPCSENPRDTEDVPWITLNSSIQKVNEWFSRSDELLGSDDSHDGESESNAKVADVLDVLNEVDEYSGSSEKIDLLASDPHEALICKSERVHSKSVESNIEDKIFGKTYRKKASLPNLSHVTENLIIGAFVTEPQIIQERPLTNKLKRKRRPTSGLHPEDFIKKADLAVQKTPEMINQGTNQTEQNGQVMNITNSGHENKTKGDSIQNEKNPNPIESLEKESAFKTKAEPISSSISNMELELNIHNSKAPKKNRLRRKSSTRHIHALELVVSRNLSPPNCTELQIDSCSSSEEIKKKKYNQMPVRHSRNLQLMEGKEPATGAKKSNKPNEQTSKRHDSDTFPELKLTNAPGSFTKCSNTSELKEFVNPSLPREEKEEKLETVKVSNNAEDPKDLMLSGERVLQTERSVESSSISLVPGTDYGTQESISLLEVSTLGKAKTEPNKCVSQCAAFENPKGLIHGCSKDNRNDTEGFKYPLGHEVNHSRETSIEMEESELDAQYLQNTFKVSKRQSFAPFSNPGNAEEECATFSAHSGSLKKQSPKVTFECEQKEENQGKNESNIKPVQTVNITAGFPVVGQKDKPVDNAKCSIKGGSRFCLSSQFRGNETGLITPNKHGLLQNPYRIPPLFPIKSFVKTKCKKNLLEENFEEHSMSPEREMGNENIPSTVSTISRNNIRENVFKEASSSNINEVGSSTNEVGSSINEIGSSDENIQAELGRNRGPKLNAMLRLGVLQPEVYKQSLPGSNCKHPEIKKQEYEEVVQTVNTDFSPYLISDNLEQPMGSSHASQVCSETPDDLLDDGEIKEDTSFAENDIKESSAVFSKSVQKGELSRSPSPFTHTHLAQGYRRGAKKLESSEENLSSEDEELPCFQHLLFGKVNNIPSQSTRHSTVATECLSKNTEENLLSLKNSLNDCSNQVILAKASQEHHLSEETKCSASLFSSQCSELEDLTANTNTQDPFLIGSSKQMRHQSESQGVGLSDKELVSDDEERGTGLEENNQEEQSMDSNLGEAASGCESETSVSEDCSGLSSQSDILTTQQRDTMQHNLIKLQQEMAELEAVLEQHGSQPSNSYPSIISDSSALEDLRNPEQSTSEKAVLTSQKSSEYPISQNPEGLSADKFEVSADSSTSKNKEPGVERSSPSKCPSLDDRWYMHSCSGSLQNRNYPSQEELIKVVDVEEQQLEESGPHDLTETSYLPRQDLEGTPYLESGISLFSDDPESDPSEDRAPESARVGNIPSSTSALKVPQLKVAESAQSPAAAHTTDTAGYNAMEESVSREKPELTASTERVNKRMSMVVSGLTPEEFMLVYKFARKHHITLTNLITEETTHVVMKTDAEFVCERTLKYFLGIAGGKWVVSYFWVTQSIKERKMLNEHDFEVRGDVVNGRNHQGPKRARESQDRKIFRGLEICCYGPFTNMPTDQLEWMVQLCGASVVKELSSFTLGTGVHPIVVVQPDAWTEDNGFHAIGQMCEAPVVTREWVLDSVALYQCQELDTYLIPQIPHSHY'],
['CALM1_HUMAN' ,'Human', 'MADQLTEEQIAEFKEAFSLFDKDGDGTITTKELGTVMRSLGQNPTEAELQDMINEVDADGNGTIDFPEFLTMMARKMKDTDSEEEIREAFRVFDKDGNGYISAAELRHVMTNLGEKLTDEEVDEMIREADIDGDGQVNYEEFVQMMTAK'],
['CCDB_ECOLI' ,'Prokaryote', 'MQFKVYTYKRESRYRLFVDVQSDIIDTPGRRMVIPLASARLLSDKVSRELYPVVHIGDESWRMMTTDMASVPVSVIGEEVADLSHRENDIKNAINLMFWGI'],
['GFP_AEQVI' ,'Other eukaryote', 'MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTLSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK'],
['GRB2_HUMAN' ,'Human', 'MEAIAKYDFKATADDELSFKRGDILKVLNEECDQNWYKAELNGKDGFIPKNYIEMKPHPWFFGKIPRAKAEEMLSKQRHDGAFLIRESESAPGDFSLSVKFGNDVQHFKVLRDGAGKYFLWVVKFNSLNELVDYHRSTSVSRNQQIFLRDIEQVPQQPTYVQALFDFDPQEDGELGFRRGDFIHVMDNSDPNWWKGACHGQTGMFPRNYVTPVNRNV'],
],
)
gr.Markdown("
")
gr.Markdown("# Fitness predictions for all single amino acid substitutions in mutation range")
gr.Markdown("Inference may take a few seconds for short proteins & mutation ranges to several minutes for longer ones")
output_image = gr.Gallery(label="Fitness predictions for all single amino acid substitutions in mutation range") #Using Gallery to break down large scoring matrices into smaller images
output_recommendations = gr.Textbox(label="Mutation recommendations")
with gr.Row():
gr.Markdown("## Download CSV Files")
output_csv_files = gr.File(label="Download CSV files with fitness scores", file_count="multiple", interactive=False)
clear_button.click(
inputs = [protein_sequence_input,mutation_range_start,mutation_range_end],
outputs = [protein_sequence_input,mutation_range_start,mutation_range_end],
fn=clear_inputs
)
run_button.click(
fn=score_and_create_matrix_all_singles,
inputs=[protein_sequence_input,mutation_range_start,mutation_range_end,model_size_selection,scoring_mirror,batch_size_inference],
outputs=[output_image,output_recommendations,output_csv_files],
)
gr.Markdown("# Mutate the starting protein sequence")
with gr.Row():
mutation_triplet = gr.Textbox(lines=1,label="Selected mutation", placeholder = "Input the mutation triplet for the selected mutation (eg., M1A)")
mutate_button = gr.Button(value="Apply mutation to starting protein", variant="primary")
mutated_protein_sequence = gr.Textbox(lines=1,label="Mutated protein sequence")
mutate_button.click(
fn = get_mutated_protein,
inputs = [protein_sequence_input,mutation_triplet],
outputs = mutated_protein_sequence
)
gr.Markdown("You may now use the output mutated sequence above as the starting sequence for another round of in silico directed evolution.
")
gr.Markdown("### About BASIS-China iGEM Team")
gr.Markdown("We are a high school synthetic biology team participating in the International Genetically Engineered Machine (iGEM) competition. Our 2025 project focuses on protein engineering and computational biology applications. This Tranception deployment is part of our broader effort to make advanced protein design tools accessible to the synthetic biology community.")
gr.Markdown("### About Tranception")
gr.Markdown("Tranception: Protein Fitness Prediction with Autoregressive Transformers and Inference-time Retrieval
Pascal Notin, Mafalda Dias, Jonathan Frazer, Javier Marchena-Hurtado, Aidan N. Gomez, Debora S. Marks*, Yarin Gal*
* equal senior authorship
")
gr.Markdown("Links: Paper Code ProteinGym BASIS-China iGEM Team")
if __name__ == "__main__":
# Don't preload models at startup - this can cause Zero GPU initialization issues
# Models will be loaded and cached on first use
print("Starting Tranception app...")
print("Note: Models will be downloaded on first use")
print("Zero GPU spaces may sleep after ~15 minutes of inactivity")
# Try to launch with ZeroGPU support first
launch_success = False
max_retries = 3
retry_count = 0
while not launch_success and retry_count < max_retries:
try:
if retry_count > 0:
print(f"Retry attempt {retry_count}/{max_retries}...")
time.sleep(2) # Wait before retry
# Launch with queue for proper Zero GPU support
tranception_design.queue(max_size=20).launch(
server_name="0.0.0.0",
server_port=7860,
show_error=True,
share=False
)
launch_success = True
except RuntimeError as e:
if "Error while initializing ZeroGPU" in str(e):
retry_count += 1
if retry_count >= max_retries:
print(f"ZeroGPU initialization failed after {max_retries} attempts")
print("Falling back to CPU mode for stability")
print("Note: The app will run slower in CPU mode")
# Update runtime mode
RUNTIME_MODE = "CPU"
# Launch without queue which doesn't trigger ZeroGPU initialization
tranception_design.launch(
server_name="0.0.0.0",
server_port=7860,
show_error=True,
share=False
)
launch_success = True
else:
# Re-raise unexpected errors
raise