Generalize thinking

app.py

@@ -73,6 +73,13 @@ LANGUAGES = [
 ACTIVE_SESSIONS = {}
 SESSION_TIMEOUT = 60
 
+THINKING_PATTERNS = [
+    # (start_token, end_token, thinking_label, answer_label)
+    ("<|begin_of_thought|>", "<|end_of_thought|>", "Thinking Process:", "Final Answer:"),
+    ("<think>", "</think>", "Thinking:", "Answer:"),
+    ("<thinking>", "</thinking>", "Thinking:", "Answer:"),
+]
+
 
 def live_count(request: gr.Request):
     current_time = time.time()
@@ -91,11 +98,29 @@ class ModelManager:
         self.model_id = None
         self.stop_generation = False
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.thinking = None  # (start, end, think_label, answer_label) or None
 
 
 model_manager = ModelManager()
 
 
+def detect_thinking(tokenizer):
+    if tokenizer is None:
+        return None
+    try:
+        special = set()
+        if hasattr(tokenizer, "additional_special_tokens"):
+            special.update(tokenizer.additional_special_tokens)
+        if hasattr(tokenizer, "added_tokens_decoder"):
+            special.update(str(v) for v in tokenizer.added_tokens_decoder.values())
+        for start, end, think_label, answer_label in THINKING_PATTERNS:
+            if start in special or start in tokenizer.get_vocab():
+                return (start, end, think_label, answer_label)
+    except Exception:
+        pass
+    return None
+
+
 class StopOnFlag(StoppingCriteria):
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
         return model_manager.stop_generation
@@ -154,7 +179,9 @@ def load_new_model(model_id):
         model_manager.tokenizer = tokenizer
         model_manager.model = model
         model_manager.model_id = model_id
-        yield f"Loaded **{model_id}** on {model_manager.device.upper()}"
+        model_manager.thinking = detect_thinking(tokenizer)
+        tag = " (thinking model)" if model_manager.thinking else ""
+        yield f"Loaded **{model_id}** on {model_manager.device.upper()}{tag}"
     except Exception as e:
         yield f"Error loading model: {str(e)}"
 
@@ -231,24 +258,25 @@ def run_inference(mode, prompt, system_prompt, context, src_lang, tgt_lang,
                   max_tokens, temperature, top_k, top_p, rep_penalty,
                   ngram_size, do_sample, gpu):
     formatted = format_prompt(mode, prompt, system_prompt, context, src_lang, tgt_lang)
+    show_prompt = mode == "Completion"
     if gpu:
         try:
-            yield from run_inference_gpu(formatted, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample)
+            yield from run_inference_gpu(formatted, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, show_prompt)
         except Exception as e:
             yield f"GPU error: {e}\n\nClick **Generate** to retry.", "GPU Unavailable"
     else:
-        yield from run_inference_raw(formatted, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample)
+        yield from run_inference_raw(formatted, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, show_prompt=show_prompt)
 
 
 @spaces.GPU(duration=estimate_duration)
-def run_inference_gpu(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample):
+def run_inference_gpu(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, show_prompt=False):
     if model_manager.model is not None:
         model_manager.model = model_manager.model.to("cuda")
 
-    yield from run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, use_cuda=True)
+    yield from run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, use_cuda=True, show_prompt=show_prompt)
 
 
-def run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, use_cuda=False):
+def run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, use_cuda=False, show_prompt=False):
     if model_manager.model is None or model_manager.tokenizer is None:
         yield "Please load a model first.", "Model not loaded"
         return
@@ -257,12 +285,10 @@ def run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_pe
 
     tokenizer = model_manager.tokenizer
     model = model_manager.model
-    model_id = model_manager.model_id
-
-    is_supra_reasoning = "Supra-50M-Reasoning" in model_id if model_id else False
+    thinking = model_manager.thinking
 
-    if is_supra_reasoning and "<|begin_of_thought|>" not in user_prompt:
-        user_prompt += "<|begin_of_thought|>\n"
+    if thinking and thinking[0] not in user_prompt:
+        user_prompt += thinking[0] + "\n"
 
     inputs = tokenizer([user_prompt], return_tensors="pt")
 
@@ -295,8 +321,11 @@ def run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_pe
     thread = Thread(target=model.generate, kwargs=generate_kwargs)
     thread.start()
 
-    if is_supra_reasoning:
-        base_display = f"Prompt: {user_prompt}\n\n{'─' * 40}\n\n"
+    if thinking:
+        base_display = ""
+        generated_text = ""
+    elif show_prompt:
+        base_display = user_prompt
         generated_text = ""
     else:
         base_display = ""
@@ -314,14 +343,21 @@ def run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_pe
 
         display_text = generated_text
 
-        if is_supra_reasoning:
+        if thinking:
+            start_tok, end_tok, think_label, answer_label = thinking
             display_text = display_text.replace("<s>", "").replace("</s>", "")
-            if not display_text.startswith("Thinking Process:"):
-                display_text = "Thinking Process:\n" + display_text
-            display_text = display_text.replace("<|begin_of_thought|>", "Thinking Process:\n")
-            display_text = display_text.replace("<|end_of_thought|>", "\n\n")
-            display_text = display_text.replace("<|begin_of_solution|>", "Final Answer:\n\n")
+            display_text = display_text.replace(start_tok, "")
+            display_text = display_text.replace(end_tok, "")
+            display_text = display_text.replace("<|begin_of_solution|>", "")
             display_text = display_text.replace("<|end_of_solution|>", "")
+            if end_tok in generated_text:
+                parts = generated_text.split(end_tok, 1)
+                think_part = parts[0].replace(start_tok, "").strip()
+                answer_part = parts[1].replace("<|begin_of_solution|>", "").replace("<|end_of_solution|>", "").strip()
+                display_text = f"> _{think_part}_\n\n**{answer_part}**"
+            else:
+                think_part = generated_text.replace(start_tok, "").replace("<s>", "").replace("</s>", "").strip()
+                display_text = f"> _{think_part}_"
 
         device_label = "CUDA" if use_cuda else "CPU"
         yield base_display + display_text, f"Speed: {tps:.2f} tokens/sec ({device_label})"
@@ -335,7 +371,17 @@ def clean_cache():
     return "Cache directory not found."
 
 
-with gr.Blocks(title="SLM Model Tester") as app:
+with gr.Blocks(title="SLM Model Tester", css="""
+#output-box {
+    background: var(--background-fill-secondary);
+    border: 1px solid var(--border-color-primary);
+    border-radius: var(--radius-lg);
+    padding: 16px;
+    min-height: 200px;
+    font-size: 15px;
+    line-height: 1.6;
+}
+""") as app:
 
     gr.Markdown("# SLM Model Evaluation Hub")
 
@@ -381,7 +427,7 @@ with gr.Blocks(title="SLM Model Tester") as app:
                 label="Prompt",
                 value="Once upon a time in a digital kingdom,",
                 placeholder="Enter your prompt here...",
-                lines=8
+                lines=3
             )
 
             context_input = gr.Textbox(
@@ -400,9 +446,7 @@ with gr.Blocks(title="SLM Model Tester") as app:
 
             run_btn = gr.Button("Generate", variant="primary", size="lg")
             status_output = gr.Markdown("*Ready*")
-            output_text = gr.Textbox(
-                label="Output", lines=15, buttons=["copy"], autoscroll=True
-            )
+            output_text = gr.Markdown(label="Output", elem_id="output-box")
 
     search_btn.click(
         fn=search_hf_models,