Update app.py

app.py

@@ -3,12 +3,21 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from transformers import TextStreamer
+import transformers
 
 import spaces
 
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained("Rorical/0-roleplay", trust_remote_code=True)
 tokenizer.chat_template = "{% for message in messages %}{{'<|im_start|>' + ((message['role'] + ':\n') if message['role'] != '' else '') + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>星野:\n' }}{% endif %}" # Be careful that this model used custom chat template.
 
+# Load model
+@spaces.GPU
+def load_pretrained_model():
+    return AutoModelForCausalLM.from_pretrained("Rorical/0-roleplay", return_dict=True, trust_remote_code=True)
+
+model = load_pretrained_model()
+
 # Define the response function
 @spaces.GPU
 def respond(
@@ -19,9 +28,6 @@ def respond(
     temperature,
     top_p,
 ):
-    # Load model and tokenizer
-    model = AutoModelForCausalLM.from_pretrained("Rorical/0-roleplay", return_dict=True, trust_remote_code=True)
-    
     # Construct the messages for the chat
     messages = [{"role": "", "content": system_message}]
     for user_message, bot_response in history: