Update app.py

app.py

@@ -1,18 +1,42 @@
 import streamlit as st
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, AutoConfig
+import json
+import os
+
+model_name = "ayyuce/NeoProtein-GPT"
+config_path = os.path.join(os.path.expanduser("~"), f".cache/huggingface/hub/models--{model_name.replace('/', '--')}/snapshots/d462becc43e0c3e4792cfa78efd029bed5dcfeb8/config.json")
+
+if not os.path.isfile(config_path):
+    config = {
+        "model_type": "gpt2",
+        "architectures": ["GPT2LMHeadModel"],
+        "vocab_size": 50257,
+        "n_positions": 1024,
+        "n_ctx": 1024,
+        "n_embd": 768,
+        "n_layer": 12,
+        "n_head": 12,
+        "activation_function": "gelu",
+        "initializer_range": 0.02,
+        "layer_norm_epsilon": 1e-5,
+        "bos_token_id": 50256,
+        "eos_token_id": 50256,
+    }
+    os.makedirs(os.path.dirname(config_path), exist_ok=True)
+    with open(config_path, "w") as f:
+        json.dump(config, f)
 
 @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, config=config, device_map="cpu")
+    model = AutoModelForCausalLM.from_pretrained(model_name, config=AutoConfig.from_pretrained(model_name))
     gen_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
     return gen_pipeline
 
 st.title("NeoProtein-GPT")
 st.write("Welcome to the NeoProtein-GPT interface. Enter a protein prompt and generate new protein sequences!")
 
+
 user_prompt = st.text_area("Enter your prompt", value="Design a novel protein sequence with a unique binding site:")
 
 if st.button("Generate Protein Sequence"):
@@ -20,4 +44,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")