Update app.py

app.py

@@ -1,22 +1,18 @@
 from flask import Flask, request, jsonify
 from transformers import pipeline
-from huggingface_hub import snapshot_download
 import torch, json, os
 
 app = Flask(__name__)
 
 MODEL_CACHE_DIR = os.environ.get("HF_HOME", "./model_cache")
 
-def ensure_model(model_name):
-    local_path = snapshot_download(model_name, cache_dir=MODEL_CACHE_DIR)
-    return local_path
-
-# Download models into cache dir
-path_120b = ensure_model("openai/gpt-oss-120b")
-path_20b  = ensure_model("openai/gpt-oss-20b")
-
-pipe_120b = pipeline("text-generation", model=path_120b, torch_dtype="auto", device_map="auto")
-pipe_20b  = pipeline("text-generation", model=path_20b, torch_dtype="auto", device_map="auto")
+# ✅ Load pre-downloaded GPT-OSS-20B model directly
+pipe_20b = pipeline(
+    "text-generation",
+    model=os.path.join(MODEL_CACHE_DIR, "models--openai--gpt-oss-20b", "snapshots"),
+    torch_dtype="auto",
+    device_map="auto"
+)
 
 MEMORY_FILE = "memory.json"
 
@@ -37,15 +33,6 @@ def generate_with_memory(pipe, session_id, user_message):
     save_memory()
     return response_text
 
-@app.route("/generate_120b")
-def generate_120b():
-    msg = request.args.get("message", "")
-    sid = request.args.get("session_id", "default")
-    if not msg.strip():
-        return jsonify({"error": "No message provided"}), 400
-    resp = generate_with_memory(pipe_120b, sid, msg)
-    return jsonify({"model": "gpt-oss-120b", "session_id": sid, "input": msg, "response": resp})
-
 @app.route("/generate_20b")
 def generate_20b():
     msg = request.args.get("message", "")