Update app.py

app.py

@@ -253,12 +253,15 @@ phi4_model = AutoModelForCausalLM.from_pretrained(
     _attn_implementation="eager",
 )
 
-# ------------------------------------------------------------------------------
-# Gradio UI configuration
-# ------------------------------------------------------------------------------
+grpo_model_name = "prithivMLmods/SmolLM2-360M-Grpo-r999"
+grpo_device = "cuda" if torch.cuda.is_available() else "cpu"
+grpo_tokenizer = AutoTokenizer.from_pretrained(grpo_model_name)
+grpo_model = AutoModelForCausalLM.from_pretrained(grpo_model_name).to(grpo_device)
+
 
 DESCRIPTION = """
-# Agent Dino 🌠"""
+# Agent Dino 🌠
+"""
 
 css = '''
 h1 {
@@ -447,7 +450,7 @@ def detect_objects(image: np.ndarray):
     
     return Image.fromarray(annotated_image)
 
-# Chat Generation Function with support for @tts, @image, @3d, @web, @rAgent, @yolo, and now @phi4 commands
+# Chat Generation Function with support for @tts, @image, @3d, @web, @rAgent, @yolo, @phi4, and now @grpo commands
 
 @spaces.GPU
 def generate(
@@ -467,7 +470,8 @@ def generate(
       - "@web": triggers a web search or webpage visit.
       - "@rAgent": initiates a reasoning chain using Llama mode.
       - "@yolo": triggers object detection using YOLO.
-      - **"@phi4": triggers multimodal (image/audio) processing using the Phi-4 model.**
+      - "@phi4": triggers multimodal (image/audio) processing using the Phi-4 model.
+      - **"@grpo": triggers text generation using the GRPO model with a text streamer.**
     """
     text = input_dict["text"]
     files = input_dict.get("files", [])
@@ -626,6 +630,37 @@ def generate(
             yield buffer
         return
 
+    # --- GRPO Text Generation branch ---
+    if text.strip().lower().startswith("@grpo"):
+        prompt = text[len("@grpo"):].strip()
+        yield "📝 Generating text with @grpo..."
+        messages = [
+            {"role": "system", "content": "Please respond in this specific format ONLY:\n<thinking>\n input your reasoning behind your answer in between these reasoning tags.\n</thinking>\n<answer>\nyour answer in between these answer tags.\n</answer>\n"},
+            {"role": "user", "content": prompt}
+        ]
+        # Use the GRPO tokenizer's chat template if available, otherwise simply join the messages.
+        input_text = grpo_tokenizer.apply_chat_template(messages, tokenize=False) if hasattr(grpo_tokenizer, "apply_chat_template") else "\n".join([msg["content"] for msg in messages])
+        inputs = grpo_tokenizer.encode(input_text, return_tensors="pt").to(grpo_model.device)
+        streamer = TextIteratorStreamer(grpo_tokenizer, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = {
+            "input_ids": inputs,
+            "max_new_tokens": 100,
+            "temperature": 0.2,
+            "top_p": 0.9,
+            "do_sample": True,
+            "use_cache": False,
+            "streamer": streamer,
+        }
+        thread = Thread(target=grpo_model.generate, kwargs=generation_kwargs)
+        thread.start()
+        buffer = ""
+        yield "🤔 Thinking..."
+        for new_text in streamer:
+            buffer += new_text
+            time.sleep(0.01)
+            yield buffer
+        return
+
     # --- Text and TTS branch ---
     tts_prefix = "@tts"
     is_tts = any(text.strip().lower().startswith(f"{tts_prefix}{i}") for i in range(1, 3))
@@ -724,6 +759,7 @@ demo = gr.ChatInterface(
         ["@rAgent Explain how a binary search algorithm works."],
         ["@web Is Grok-3 Beats DeepSeek-R1 at Reasoning ?"],
         ["@tts1 Explain Tower of Hanoi"],
+        ["@grpo If there are 12 cookies in a dozen and you have 5 dozen, how many cookies do you have?"],
     ],
     cache_examples=False,
     type="messages",
@@ -734,7 +770,7 @@ demo = gr.ChatInterface(
         label="Query Input", 
         file_types=["image", "audio"],
         file_count="multiple", 
-        placeholder="@tts1, @tts2, @image, @3d, @phi4 [image, audio], @rAgent, @web, @yolo, default [plain text]"
+        placeholder="@tts1, @tts2, @image, @3d, @phi4 [image, audio], @rAgent, @web, @yolo, @grpo, default [plain text]"
     ),
     stop_btn="Stop Generation",
     multimodal=True,