import torch
from transformers import AutoTokenizer, AutoModel
import gradio as gr

# Load model & tokenizer once at startup
MODEL_NAME = "facebook/esm2_t33_650M_UR50D"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME)

# Define your inference logic
def run_esm2(sequence: str):
    # Tokenize input
    inputs = tokenizer(sequence, return_tensors="pt")
    
    # Run model
    with torch.no_grad():
        outputs = model(**inputs)
    
    # Return last hidden state (or modify to return something else)
    return outputs.last_hidden_state[0].tolist()  # [[float, float, ...], ...]

# Gradio interface
demo = gr.Interface(
    fn=run_esm2,
    inputs=gr.Textbox(label="Protein Sequence"),
    outputs="json",
    title="ESM-2 Inference API",
    description="Paste a protein sequence to get the output embeddings.",
)

# Launch app
demo.launch()