Update app.py

app.py

@@ -1,12 +1,12 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import torch
 
 st.set_page_config(page_title="NeoProtein Designer", page_icon="🧬")
 st.title("🧬 NeoProtein-GPT Protein Designer")
 st.markdown("""
 ### Design novel protein sequences with unique binding sites
-*Using the [NeoProtein-GPT](https://huggingface.co/ayyuce/NeoProtein-GPT) model from Hugging Face*
+*Using the [NeoProtein-GPT](https://huggingface.co/ayyuce/NeoProtein-GPT) model*
 """)
 
 with st.sidebar:
@@ -18,12 +18,14 @@ with st.sidebar:
 
 @st.cache_resource(show_spinner=False)
 def load_model():
-    model = AutoModelForCausalLM.from_pretrained(
+    model = GPT2LMHeadModel.from_pretrained(
         "ayyuce/NeoProtein-GPT",
-        device_map="cpu",
-        torch_dtype=torch.float32
+        force_download=True,
+        resume_download=False,
+        local_files_only=False,
+        trust_remote_code=True
     )
-    tokenizer = AutoTokenizer.from_pretrained("ayyuce/NeoProtein-GPT")
+    tokenizer = GPT2Tokenizer.from_pretrained("ayyuce/NeoProtein-GPT")
     return model, tokenizer
 
 model, tokenizer = load_model()
@@ -61,6 +63,7 @@ def generate_sequences():
         st.error(f"Generation failed: {str(e)}")
         return []
 
+
 if st.button("Generate Protein Sequences"):
     with st.spinner("Designing novel proteins..."):
         sequences = generate_sequences()
@@ -74,7 +77,6 @@ if st.button("Generate Protein Sequences"):
             {seq}
             ```
             """)
-            st.divider()
 
 st.markdown("""
 ### How to use: