Update app.py

app.py

@@ -79,18 +79,28 @@ def load_pipeline(model_name):
     Load and cache a transformers pipeline for text generation.
     Tries bfloat16, falls back to float16 or float32 if unsupported.
     """
-    global PIPELINES
+    
     if model_name in PIPELINES.keys():
         return PIPELINES[model_name]
     repo = MODELS[model_name]["repo_id"]
-    tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-            repo,
-            device_map=device,
-            trust_remote_code=True,
-        )
+    if model_name == "secgpt-mini":
+        tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True, subfolder="models")
+        model = AutoModelForCausalLM.from_pretrained(
+                repo,
+                device_map=device,
+                trust_remote_code=True,
+                subfolder="models",
+            )
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
+        model = AutoModelForCausalLM.from_pretrained(
+                repo,
+                device_map=device,
+                trust_remote_code=True,
+            )
+    global PIPELINES
     PIPELINES[model_name] = {"tokenizer": tokenizer, "model": model}
-    return PIPELINES[model_name]
+    return {"tokenizer": tokenizer, "model": model}
 
 
 def retrieve_context(query, max_results=6, max_chars=600):