refactor: update llm_zerogpu.py and llm_modal to align each other

src/unpredictable_lord/llm_modal.py

@@ -18,8 +18,11 @@ MOUNT_DIR = "/data"
 
 # https://huggingface.co/openai/gpt-oss-20b
 MODEL_IDENTIFIER = "openai/gpt-oss-20b"
-# https://huggingface.co/openai/gpt-oss-120b
-# MODEL_IDENTIFIER = "openai/gpt-oss-120b"
+
+# MAX_MODEL_TOKENS >= Input + Output
+MAX_MODEL_TOKENS = 64 * 1024  # gpt-oss models support up to 128k(128*1024) tokens
+MAX_OUTPUT_TOKENS = 512
+
 
 # https://modal.com/docs/guide/gpu#specifying-gpu-type
 GPU_NAME = "L4"
@@ -39,10 +42,6 @@ GPU = f"{GPU_NAME}:{GPU_NUM}"
 # | L4        |  24 GB | $0.80 /h |
 # | T4        |  16 GB | $0.59 /h |
 
-# MAX_MODEL_TOKENS >= Input + Output
-MAX_MODEL_TOKENS = 64 * 1024  # gpt-oss models support up to 128k(128*1024) tokens
-MAX_OUTPUT_TOKENS = 512
-
 image = (
     # https://hub.docker.com/r/nvidia/cuda/tags?name=12.8
     # https://hub.docker.com/layers/nvidia/cuda/12.8.1-devel-ubuntu24.04
@@ -104,9 +103,6 @@ def load_model():
 @app.function(
     gpu=GPU,
     volumes={MOUNT_DIR: MOUNT_VOLUME},
-    # secrets=[modal.Secret.from_name("huggingface-secret")],
-    # scaledown_window=15 * 60,
-    # timeout=30 * 60,
 )
 def generate_stream(input_tokens):
     """

src/unpredictable_lord/llm_zerogpu.py

@@ -1,3 +1,5 @@
+import logging
+import subprocess
 from threading import Thread
 
 import openai_harmony as oh
@@ -7,27 +9,45 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from unpredictable_lord.tokenstreamer import TokenStreamer
 
+logger = logging.getLogger(__name__)
+
 # https://huggingface.co/openai/gpt-oss-20b
 MODEL_IDENTIFIER = "openai/gpt-oss-20b"
-# https://huggingface.co/openai/gpt-oss-120b
-# MODEL_IDENTIFIER = "openai/gpt-oss-120b"
 
 # MAX_MODEL_TOKENS >= Input + Output
 MAX_MODEL_TOKENS = 64 * 1024  # gpt-oss models support up to 128k(128*1024) tokens
 MAX_OUTPUT_TOKENS = 512
 
 
-# Global model and tokenizer (loaded once at module import)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_IDENTIFIER)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_IDENTIFIER,
-    torch_dtype="auto",
-    device_map="auto",
-)
+# Global model and tokenizer (loaded once)
+model = None
+tokenizer = None
+stop_token_ids = None
+
+
+def load_model():
+    """Load model and tokenizer into global variables."""
+    global model, tokenizer, stop_token_ids
+
+    if model is not None:
+        return
+
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_IDENTIFIER)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_IDENTIFIER,
+        torch_dtype="auto",
+        device_map="auto",
+    )
+
+    # Load stop token IDs
+    _encoding = oh.load_harmony_encoding(oh.HarmonyEncodingName.HARMONY_GPT_OSS)
+    stop_token_ids = _encoding.stop_tokens_for_assistant_actions()
+
+    # Show GPU information
+    subprocess.run(["nvidia-smi"])
+
 
-# Load stop token IDs
-encoding = oh.load_harmony_encoding(oh.HarmonyEncodingName.HARMONY_GPT_OSS)
-stop_token_ids = encoding.stop_tokens_for_assistant_actions()
+load_model()
 
 
 @spaces.GPU