Update app.py (#1)

- Update app.py (81e19ee3e2e3b367ae549ae4084774914687aa09)

app.py

@@ -14,7 +14,7 @@ from threading import Thread
 
 
 # =============================================================================
-# CONFIG
+# ROSTAM R1 CONFIG
 # =============================================================================
 
 MODEL_ID = "prog-love/rostam-r1-stage13-1"
@@ -22,38 +22,77 @@ MODEL_ID = "prog-love/rostam-r1-stage13-1"
 print("=" * 70)
 print("ROSTAM R1")
 print("=" * 70)
-print(f"Loading model: {MODEL_ID}")
-print(f"Transformers: {transformers.__version__}")
-print(f"PyTorch: {torch.__version__}")
-print(f"CUDA available: {torch.cuda.is_available()}")
+print(f"Model        : {MODEL_ID}")
+print(f"Transformers : {transformers.__version__}")
+print(f"PyTorch      : {torch.__version__}")
+print(f"CUDA         : {torch.cuda.is_available()}")
 
 if torch.cuda.is_available():
-    print(f"CUDA device: {torch.cuda.get_device_name(0)}")
+    print(f"GPU          : {torch.cuda.get_device_name(0)}")
 
 
 # =============================================================================
-# GEMMA 4 / HETEROGENEOUS CONFIG COMPATIBILITY
+# GEMMA 4 / TRANSFORMERS 5.14.1 COMPATIBILITY FIX
 # =============================================================================
 #
-# The current Transformers stack can reject Gemma 4 configs containing
-# per-layer attributes such as `head_dim` when architecture validation tries
-# to access them globally.
+# Transformers 5.x supports heterogeneous/per-layer configurations.
 #
-# We explicitly allow that access while the configuration is being created.
+# Rostam's Gemma 4 config contains per-layer attributes such as `head_dim`.
+# During Gemma4TextConfig.__post_init__(), Transformers performs architecture
+# validation and accesses config.head_dim globally.
 #
-# This does NOT modify the model checkpoint or its weights.
+# In a heterogeneous config this normally raises:
+#
+#   AmbiguousGlobalPerLayerAttributeError
+#
+# The important detail is that the compatibility flag MUST exist BEFORE
+# Gemma4TextConfig.__post_init__() runs.
+#
+# Setting it after AutoConfig.from_pretrained() is too late.
+#
+# Therefore we wrap Gemma4TextConfig.__post_init__ and inject the flag before
+# the original validation code executes.
 # =============================================================================
 
+print("Installing Gemma 4 configuration compatibility patch...")
+
 try:
-    from transformers.integrations.heterogeneity.configuration_utils import (
-        AmbiguousGlobalPerLayerAttributeError,
+
+    from transformers.models.gemma4.configuration_gemma4 import (
+        Gemma4TextConfig,
     )
 
-    print("Detected Transformers heterogeneous configuration support.")
+    _original_gemma4_text_post_init = Gemma4TextConfig.__post_init__
+
+    def _rostam_gemma4_text_post_init(self, *args, **kwargs):
+
+        # Enable intentional global access to per-layer attributes.
+        #
+        # The model architecture itself still retains the per-layer
+        # configuration. This only allows Transformers' architecture
+        # validation to read the global fallback value.
+        object.__setattr__(
+            self,
+            "allow_global_per_layer_attribute_access",
+            True,
+        )
+
+        return _original_gemma4_text_post_init(
+            self,
+            *args,
+            **kwargs,
+        )
+
+    Gemma4TextConfig.__post_init__ = _rostam_gemma4_text_post_init
 
-except ImportError:
-    AmbiguousGlobalPerLayerAttributeError = None
-    print("No heterogeneous configuration module detected.")
+    print("Gemma 4 compatibility patch installed.")
+
+except Exception as e:
+
+    print(
+        "WARNING: Could not install Gemma 4 compatibility patch:"
+    )
+    print(repr(e))
 
 
 # =============================================================================
@@ -67,44 +106,37 @@ config = AutoConfig.from_pretrained(
     trust_remote_code=True,
 )
 
-# The exact flag used by the error message.
-#
-# It may live directly on the configuration object depending on the
-# Transformers version.
-if hasattr(config, "allow_global_per_layer_attribute_access"):
-    config.allow_global_per_layer_attribute_access = True
+print("Configuration loaded successfully.")
 
-# Gemma4 contains a nested text configuration.
-if hasattr(config, "text_config"):
-    try:
-        config.text_config.allow_global_per_layer_attribute_access = True
-    except Exception:
-        pass
 
+# =============================================================================
+# SAFETY: ENSURE GLOBAL ACCESS IS ENABLED ON FINAL CONFIG
+# =============================================================================
 
-# Some Transformers versions expose the flag through the heterogeneity
-# configuration machinery instead of the top-level config.
-#
-# Set it defensively where available.
-for cfg in (
-    config,
-    getattr(config, "text_config", None),
-):
+try:
+
+    object.__setattr__(
+        config,
+        "allow_global_per_layer_attribute_access",
+        True,
+    )
+
+except Exception:
+    pass
+
+
+try:
 
-    if cfg is None:
-        continue
+    if hasattr(config, "text_config"):
 
-    try:
         object.__setattr__(
-            cfg,
+            config.text_config,
             "allow_global_per_layer_attribute_access",
             True,
         )
-    except Exception:
-        pass
-
 
-print("Configuration loaded.")
+except Exception:
+    pass
 
 
 # =============================================================================
@@ -118,7 +150,7 @@ tokenizer = AutoTokenizer.from_pretrained(
     trust_remote_code=True,
 )
 
-print("Tokenizer loaded.")
+print("Tokenizer loaded successfully.")
 
 
 # =============================================================================
@@ -139,9 +171,15 @@ model.eval()
 
 print("Model loaded successfully.")
 
+
+# =============================================================================
+# GPU INFORMATION
+# =============================================================================
+
 if torch.cuda.is_available():
-    allocated = torch.cuda.memory_allocated() / 1024**3
-    reserved = torch.cuda.memory_reserved() / 1024**3
+
+    allocated = torch.cuda.memory_allocated() / (1024 ** 3)
+    reserved = torch.cuda.memory_reserved() / (1024 ** 3)
 
     print(
         f"GPU memory: "
@@ -149,15 +187,22 @@ if torch.cuda.is_available():
         f"{reserved:.2f} GB reserved"
     )
 
+
+print("=" * 70)
+print("ROSTAM IS READY")
 print("=" * 70)
 
 
 # =============================================================================
-# CSS
+# UI CSS
 # =============================================================================
 
 CUSTOM_CSS = """
-/* Persian / RTL messages */
+
+/* -------------------------------------------------------------------------
+   Persian / RTL chat
+   ------------------------------------------------------------------------- */
+
 .message-wrap {
     direction: rtl !important;
     text-align: right !important;
@@ -168,34 +213,50 @@ CUSTOM_CSS = """
     text-align: right !important;
 }
 
-/* Chat bubbles */
 .chatbot {
-    direction: rtl;
+    direction: rtl !important;
 }
 
-/* User input */
+
+/* -------------------------------------------------------------------------
+   Persian input
+   ------------------------------------------------------------------------- */
+
 textarea {
     direction: rtl !important;
     text-align: right !important;
 }
 
-/* Hide Gradio footer */
-footer {
-    display: none !important;
-}
 
-/* Better Persian rendering */
+/* -------------------------------------------------------------------------
+   Markdown
+   ------------------------------------------------------------------------- */
+
 .prose {
     direction: rtl;
     text-align: right;
 }
 
-/* Keep code blocks LTR */
+
+/* -------------------------------------------------------------------------
+   Code remains LTR
+   ------------------------------------------------------------------------- */
+
 pre,
 code {
     direction: ltr !important;
     text-align: left !important;
 }
+
+
+/* -------------------------------------------------------------------------
+   Hide Gradio footer
+   ------------------------------------------------------------------------- */
+
+footer {
+    display: none !important;
+}
+
 """
 
 
@@ -212,15 +273,13 @@ def generate_response(
     top_p,
     max_new_tokens,
 ):
-    """
-    Generate a streaming response from Rostam.
-    """
 
     messages = []
 
-    # -------------------------------------------------------------------------
-    # System prompt
-    # -------------------------------------------------------------------------
+
+    # =========================================================================
+    # SYSTEM PROMPT
+    # =========================================================================
 
     if system_prompt and system_prompt.strip():
 
@@ -231,9 +290,10 @@ def generate_response(
             }
         )
 
-    # -------------------------------------------------------------------------
-    # Conversation history
-    # -------------------------------------------------------------------------
+
+    # =========================================================================
+    # HISTORY
+    # =========================================================================
 
     if history:
 
@@ -257,9 +317,10 @@ def generate_response(
                     }
                 )
 
-    # -------------------------------------------------------------------------
-    # Current message
-    # -------------------------------------------------------------------------
+
+    # =========================================================================
+    # CURRENT USER MESSAGE
+    # =========================================================================
 
     messages.append(
         {
@@ -268,9 +329,10 @@ def generate_response(
         }
     )
 
-    # -------------------------------------------------------------------------
-    # Chat template
-    # -------------------------------------------------------------------------
+
+    # =========================================================================
+    # CHAT TEMPLATE
+    # =========================================================================
 
     prompt = tokenizer.apply_chat_template(
         messages,
@@ -278,18 +340,21 @@ def generate_response(
         add_generation_prompt=True,
     )
 
-    # -------------------------------------------------------------------------
-    # Tokenization
-    # -------------------------------------------------------------------------
+
+    # =========================================================================
+    # TOKENIZE
+    # =========================================================================
 
     inputs = tokenizer(
         prompt,
         return_tensors="pt",
     )
 
-    # device_map="auto" places the model on the appropriate device.
-    #
-    # For a single-GPU deployment this normally means CUDA.
+
+    # =========================================================================
+    # MOVE INPUTS TO MODEL DEVICE
+    # =========================================================================
+
     if torch.cuda.is_available():
 
         inputs = {
@@ -297,9 +362,10 @@ def generate_response(
             for key, value in inputs.items()
         }
 
-    # -------------------------------------------------------------------------
-    # Streamer
-    # -------------------------------------------------------------------------
+
+    # =========================================================================
+    # STREAMER
+    # =========================================================================
 
     streamer = TextIteratorStreamer(
         tokenizer,
@@ -307,14 +373,16 @@ def generate_response(
         skip_special_tokens=True,
     )
 
-    # -------------------------------------------------------------------------
-    # Generation parameters
-    # -------------------------------------------------------------------------
+
+    # =========================================================================
+    # GENERATION SETTINGS
+    # =========================================================================
 
     temperature = float(temperature)
     top_p = float(top_p)
     max_new_tokens = int(max_new_tokens)
 
+
     generation_kwargs = {
         **inputs,
 
@@ -331,9 +399,10 @@ def generate_response(
         ),
     }
 
-    # -------------------------------------------------------------------------
-    # Sampling
-    # -------------------------------------------------------------------------
+
+    # =========================================================================
+    # SAMPLING
+    # =========================================================================
 
     if temperature <= 0:
 
@@ -345,9 +414,10 @@ def generate_response(
         generation_kwargs["temperature"] = temperature
         generation_kwargs["top_p"] = top_p
 
-    # -------------------------------------------------------------------------
-    # Start generation thread
-    # -------------------------------------------------------------------------
+
+    # =========================================================================
+    # GENERATION THREAD
+    # =========================================================================
 
     thread = Thread(
         target=model.generate,
@@ -357,9 +427,10 @@ def generate_response(
 
     thread.start()
 
-    # -------------------------------------------------------------------------
-    # Stream output
-    # -------------------------------------------------------------------------
+
+    # =========================================================================
+    # STREAM OUTPUT
+    # =========================================================================
 
     partial_text = ""
 
@@ -371,7 +442,7 @@ def generate_response(
 
 
 # =============================================================================
-# GRADIO UI
+# GRADIO APPLICATION
 # =============================================================================
 
 with gr.Blocks(
@@ -380,9 +451,10 @@ with gr.Blocks(
     title="رستم",
 ) as demo:
 
-    # -------------------------------------------------------------------------
-    # Header
-    # -------------------------------------------------------------------------
+
+    # =========================================================================
+    # HEADER
+    # =========================================================================
 
     gr.Markdown(
         """
@@ -392,9 +464,10 @@ with gr.Blocks(
         """
     )
 
-    # -------------------------------------------------------------------------
-    # Advanced settings
-    # -------------------------------------------------------------------------
+
+    # =========================================================================
+    # ADVANCED SETTINGS
+    # =========================================================================
 
     with gr.Accordion(
         "تنظیمات پیشرفته",
@@ -407,6 +480,7 @@ with gr.Blocks(
             lines=2,
         )
 
+
         with gr.Row():
 
             temperature = gr.Slider(
@@ -417,6 +491,7 @@ with gr.Blocks(
                 label="Temperature",
             )
 
+
             top_p = gr.Slider(
                 minimum=0.1,
                 maximum=1.0,
@@ -425,6 +500,7 @@ with gr.Blocks(
                 label="Top-p",
             )
 
+
             max_new_tokens = gr.Slider(
                 minimum=32,
                 maximum=8000,
@@ -433,9 +509,10 @@ with gr.Blocks(
                 label="حداکثر طول پاسخ",
             )
 
-    # -------------------------------------------------------------------------
-    # Chat
-    # -------------------------------------------------------------------------
+
+    # =========================================================================
+    # CHAT
+    # =========================================================================
 
     chat = gr.ChatInterface(
         fn=generate_response,
@@ -455,6 +532,7 @@ with gr.Blocks(
                 0.9,
                 512,
             ],
+
             [
                 "تو چه مدلی هستی؟",
                 "",
@@ -462,6 +540,7 @@ with gr.Blocks(
                 0.9,
                 512,
             ],
+
             [
                 "دربارهٔ رستم و سهراب برام بگو",
                 "",