Update app.py

app.py

@@ -25,6 +25,12 @@ if not torch.cuda.is_available():
 
 models_cache = {}
 
+def load_model(model_id: str):
+    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_4bit=True)
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    tokenizer.use_default_system_prompt = False
+    return model, tokenizer
+
 @spaces.GPU(duration=120)
 def generate(
     model_id: str,
@@ -37,14 +43,12 @@ def generate(
     top_k: int = 50,
     repetition_penalty: float = 1.2,
 ) -> Iterator[str]:
+    # Load the model if it's not already loaded
     if model_id not in models_cache:
-        model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_4bit=True)
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        tokenizer.use_default_system_prompt = False
+        model, tokenizer = load_model(model_id)
         models_cache[model_id] = (model, tokenizer)
     else:
         model, tokenizer = models_cache[model_id]
-
     conversation = []
     if system_prompt:
         conversation.append({"role": "system", "content": system_prompt})