ProstT5：蛋白质序列与结构的双语语言模型¶

# 安装依赖
pip install torch
pip install transformers
pip install sentencepiece

# 针对最新 transformers 版本可能出现的 protobuf 错误
pip install protobuf

from transformers import T5Tokenizer, T5EncoderModel
import torch
import re

# 选择设备
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# 加载分词器与模型（编码器模式）
tokenizer = T5Tokenizer.from_pretrained('Rostlab/ProstT5', do_lower_case=False)
model = T5EncoderModel.from_pretrained("Rostlab/ProstT5").to(device)

# GPU 使用半精度加速，CPU 则保持全精度
if device.type == 'cpu':
    model.float()
else:
    model.half()

# 准备序列：氨基酸大写，3Di 小写
sequence_examples = ["PRTEINO", "strct"]

# 替换稀有氨基酸为 X，并在字符间插入空格
sequence_examples = [" ".join(list(re.sub(r"[UZOB]", "X", seq))) for seq in sequence_examples]

# 添加方向前缀：大写序列用 <AA2fold>，小写结构用 <fold2AA>
prefixed_sequences = []
for seq in sequence_examples:
    if seq.isupper():
        prefixed_sequences.append("<AA2fold> " + seq)
    else:
        prefixed_sequences.append("<fold2AA> " + seq)

# 批编码并补齐
encodings = tokenizer.batch_encode_plus(
    prefixed_sequences,
    add_special_tokens=True,
    padding="longest",
    return_tensors='pt'
).to(device)

# 前向传播提取嵌入
with torch.no_grad():
    embedding_result = model(
        input_ids=encodings.input_ids,
        attention_mask=encodings.attention_mask
    )

# 获取残基级别嵌入（跳过特殊 token 和前缀 token）
# 第一条序列："PRTEINO" 长度 7，前缀 "<AA2fold>" 占用 1 个 token，加上 [CLS] 等特殊 token（默认无 CLS，需确认）
# 实际 token 序列：<AA2fold> + 空格分割后的字符，共 8 个 token，索引 0 为 <AA2fold>
emb_0 = embedding_result.last_hidden_state[0, 1:8]  # 形状 (7, 1024)
emb_1 = embedding_result.last_hidden_state[1, 1:6]  # 形状 (5, 1024)

# 计算蛋白质级别的全局嵌入（对残基取平均）
protein_emb_0 = emb_0.mean(dim=0)  # 形状 (1024,)

from transformers import T5Tokenizer, AutoModelForSeq2SeqLM
import torch
import re

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# 加载分词器与模型（Seq2Seq 模式）
tokenizer = T5Tokenizer.from_pretrained('Rostlab/ProstT5', do_lower_case=False)
model = AutoModelForSeq2SeqLM.from_pretrained("Rostlab/ProstT5").to(device)
model = model.float() if device.type == 'cpu' else model.half()

# 输入氨基酸序列（大写）
sequence_examples = ["PRTEINO", "SEQWENCE"]
min_len = min(len(s) for s in sequence_examples)
max_len = max(len(s) for s in sequence_examples)

# 预处理：替换稀有氨基酸、插入空格
sequence_examples = [" ".join(list(re.sub(r"[UZOB]", "X", seq))) for seq in sequence_examples]

# 添加 AA -> 3Di 前缀
prefixed_sequences = ["<AA2fold> " + s for s in sequence_examples]

# 编码
encodings = tokenizer.batch_encode_plus(
    prefixed_sequences,
    add_special_tokens=True,
    padding="longest",
    return_tensors='pt'
).to(device)

# 生成参数（AA -> 3Di）
gen_kwargs_aa2fold = {
    "do_sample": True,
    "num_beams": 3,
    "top_p": 0.95,
    "temperature": 1.2,
    "top_k": 6,
    "repetition_penalty": 1.2,
}

# 生成翻译结果
with torch.no_grad():
    translations = model.generate(
        input_ids=encodings.input_ids,
        attention_mask=encodings.attention_mask,
        max_length=max_len,
        min_length=min_len,
        early_stopping=True,
        num_return_sequences=1,
        **gen_kwargs_aa2fold
    )

# 解码并去除空格，恢复小写 3Di 字符串
decoded_texts = tokenizer.batch_decode(translations, skip_special_tokens=True)
structure_sequences = ["".join(ts.split()) for ts in decoded_texts]

# ---------- 反向翻译：3Di -> AA ----------
# 准备反向输入（加上 <fold2AA> 前缀）
back_prefixed = ["<fold2AA> " + ts for ts in decoded_texts]  # ts 依然保留空格分隔
back_encodings = tokenizer.batch_encode_plus(
    back_prefixed,
    add_special_tokens=True,
    padding="longest",
    return_tensors='pt'
).to(device)

# 生成氨基酸序列（也可使用不同的生成参数）
with torch.no_grad():
    back_translations = model.generate(
        input_ids=back_encodings.input_ids,
        attention_mask=back_encodings.attention_mask,
        max_length=max_len,
        min_length=min_len,
        early_stopping=True,
        num_return_sequences=1,
        do_sample=True,
        top_p=0.95,
        temperature=1.0,
        repetition_penalty=1.2,
    )

back_texts = tokenizer.batch_decode(back_translations, skip_special_tokens=True)
aa_sequences = ["".join(ts.split()) for ts in back_texts]  # 最终氨基酸序列（大写）

model = T5EncoderModel.from_pretrained("Rostlab/ProstT5")
# 输入： "<AA2fold> P R T E I N O" 或 "<fold2AA> s t r c t"
embeddings = model(**encoded_inputs).last_hidden_state

model = AutoModelForSeq2SeqLM.from_pretrained("Rostlab/ProstT5")
# 前向翻译
input_ids = tokenizer("<AA2fold> P R T E I N O", return_tensors="pt").input_ids
output_ids = model.generate(input_ids, max_length=length)
structure = tokenizer.decode(output_ids[0], skip_special_tokens=True)
# 反向翻译
input_ids = tokenizer("<fold2AA> s t r c t", return_tensors="pt").input_ids
output_ids = model.generate(input_ids, max_length=length)
sequence = tokenizer.decode(output_ids[0], skip_special_tokens=True)