Update app.py

app.py

@@ -24,6 +24,14 @@ from transformers import (
 from transformers.image_utils import load_image
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
+# Load the reasoning model interface from sambanova_gradio
+try:
+    import sambanova_gradio
+    reasoning_interface = gr.load("DeepSeek-R1-Distill-Llama-70B", src=sambanova_gradio.registry, accept_token=True)
+except Exception as e:
+    reasoning_interface = None
+    print("Reasoning model could not be loaded:", e)
+
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
@@ -186,8 +194,8 @@ def generate(
 ):
     text = input_dict["text"]
     files = input_dict.get("files", [])
-
     lower_text = text.lower().strip()
+
     # Check if the prompt is an image generation command using model flags.
     if (lower_text.startswith("@lightningv5") or 
         lower_text.startswith("@lightningv4") or 
@@ -240,6 +248,18 @@ def generate(
         yield gr.Image(image_path)
         return
 
+    # New reasoning feature.
+    elif lower_text.startswith("@reasoning"):
+        prompt_clean = re.sub(r"@reasoning", "", text, flags=re.IGNORECASE).strip()
+        if reasoning_interface is None:
+            yield "Reasoning model is not available."
+            return
+        yield "Reasoning in progress..."
+        # Call the reasoning model's prediction.
+        result = reasoning_interface.predict(prompt_clean)
+        yield result
+        return
+
     # Otherwise, handle text/chat (and TTS) generation.
     tts_prefix = "@tts"
     is_tts = any(text.strip().lower().startswith(f"{tts_prefix}{i}") for i in range(1, 3))
@@ -345,6 +365,7 @@ demo = gr.ChatInterface(
         ['@turbov3 "Abstract art, colorful and vibrant"'],
         ["Write a Python function to check if a number is prime."],
         ["@tts2 What causes rainbows to form?"],
+        ["@reasoning Explain the significance of Gödel's incompleteness theorems."],
     ],
     cache_examples=False,
     type="messages",