Update app.py

app.py

@@ -1,11 +1,12 @@
 import streamlit as st
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, pipeline
 
 @st.cache_resource(show_spinner=False)
 def load_generator():
     model_name = "ayyuce/NeoProtein-GPT"
+    config = AutoConfig.from_pretrained(model_name, model_type="gpt2")
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu")
+    model = AutoModelForCausalLM.from_pretrained(model_name, config=config, device_map="cpu")
     gen_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
     return gen_pipeline
 
@@ -19,4 +20,4 @@ if st.button("Generate Protein Sequence"):
         outputs = load_generator()(user_prompt, max_new_tokens=200, do_sample=True, temperature=0.7)
         generated_text = outputs[0]["generated_text"]
     st.subheader("Generated Sequence:")
-    st.code(generated_text, language="python")
+    st.code(generated_text, language="python")