import streamlit as st
from transformers import AutoTokenizer, EsmModel
import torch

model_name = "facebook/esm2_t6_8M_UR50D"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = EsmModel.from_pretrained(model_name)

aa_seq = st.text_input('Type AA sequance here')

inputs = tokenizer(aa_seq, return_tensors="pt")
outputs = model(**inputs)

last_hidden_states = outputs.last_hidden_state

st.write(model_name, 'last hidden states',last_hidden_states.shape,  last_hidden_states, 'Also, Dania is not gay')