Update app.py

app.py

@@ -9,6 +9,12 @@ model = AutoModelForCausalLM.from_pretrained("unsloth/gemma-3-1b-it")
 model.load_adapter("Oysiyl/gemma-3-1B-GRPO")
 tokenizer = AutoTokenizer.from_pretrained("unsloth/gemma-3-1b-it")
 
+if torch.cuda.is_available():
+    model.to("cuda")
+
+if torch.backends.mps.is_available():
+    model.to("mps")
+
 
 def process_history(history):
     """Process chat history into the format expected by the model."""
@@ -27,7 +33,7 @@ def process_new_user_message(message):
 
 
 def respond(
-    message,
+    user_message,
     history: list[tuple[str, str]],
     system_message,
     max_tokens,
@@ -37,27 +43,26 @@ def respond(
     # Format messages according to Gemma's expected chat format
     messages = []
     if system_message:
-        messages.append({"role": "system", "content": [{"type": "text", "text": system_message}]})
+        messages.append({"role": "system", "content": system_message})
     
     # Process the conversation history
     if history:
         messages.extend(process_history(history))
     
     # Add the new user message
-    messages.append({"role": "user", "content": process_new_user_message(message)})
+    messages.append({"role": "user", "content": user_message})
     
     # Apply chat template
-    inputs = tokenizer.apply_chat_template(
+    prompt = tokenizer.apply_chat_template(
         messages,
         add_generation_prompt=True,
-        tokenize=True,
-        return_dict=True,
-        return_tensors="pt"
+        tokenize=False,
     )
-    
+    inputs = tokenizer(prompt, return_tensors="pt")
     if torch.cuda.is_available():
         inputs = inputs.to("cuda")
-        model.to("cuda")
+    elif torch.backends.mps.is_available():
+        inputs = inputs.to("mps")
     
     # Set up the streamer
     streamer = TextIteratorStreamer(tokenizer, timeout=30.0, skip_prompt=True, skip_special_tokens=False)
@@ -70,31 +75,50 @@ def respond(
         temperature=temperature,
         top_p=top_p,
         top_k=64,  # Recommended Gemma-3 setting
-        do_sample=True,
+        # do_sample=True,
     )
     
     thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
     thread.start()
     
-    # Stream the output, add "Thinking" at the beginning
-    output = "Thinking: \n"
+    output = ""
     for token in streamer:
         output += token
-        # Check if "<SOLUTION>" token is in the output and format everything after it as bold
-        if "<SOLUTION>" in output:
-            solution_start = output.find("<SOLUTION>") + len("<SOLUTION>")
-            solution_end = output.find("</SOLUTION>")
-            if solution_end > solution_start:
-                formatted_output = (
-                    output[:solution_start] + 
-                    "Final answer: **" + output[solution_start:solution_end] + "**" + 
-                    output[solution_end:]
-                )
-                yield formatted_output
+
+        # Tags
+        start_tag = "<start_working_out>"
+        sol_start = "<SOLUTION>"
+
+        thinking = ""
+        final_answer = ""
+
+        # Extract "Thinking" section: everything after <start_working_out>
+        if start_tag in output:
+            start_idx = output.find(start_tag) + len(start_tag)
+            # If <SOLUTION> is also present, stop "Thinking" at <SOLUTION>
+            if sol_start in output:
+                end_idx = output.find(sol_start)
             else:
-                yield output
-        else:
-            yield output
+                end_idx = len(output)
+            thinking = output[start_idx:end_idx].strip()
+
+        # Extract "Final answer" section: everything after <SOLUTION>
+        if sol_start in output:
+            sol_start_idx = output.find(sol_start) + len(sol_start)
+            final_answer = output[sol_start_idx:].strip()
+
+        # Build formatted output
+        formatted_output = ""
+        if thinking:
+            formatted_output += "### Thinking:\n" + thinking + "\n"
+        if final_answer:
+            formatted_output += "\n### Final answer:\n**" + final_answer + "**"
+
+        # If nothing found yet, just show the raw output (for streaming effect)
+        if not thinking and not final_answer:
+            formatted_output = output
+
+        yield formatted_output
 
 
 """
@@ -118,10 +142,11 @@ demo = gr.ChatInterface(
         ),
     ],
     examples=[
-        ["Apple sold 100 iPhones at their New York store today for an average cost of $1000. They also sold 20 iPads for an average cost of $900 and 80 Apple TVs for an average cost of $200. What was the average cost across all products sold today?"],
-        ["Julie is reading a 120-page book. Yesterday, she was able to read 12 pages and today, she read twice as many pages as yesterday. If she wants to read half of the remaining pages tomorrow, how many pages should she read?"],
-        ["Mel is three years younger than Katherine. When Katherine is two dozen years old, how old will Mel be in years?"],
-        ["What is the sqrt of 101?"],
+        ["Toulouse has twice as many sheep as Charleston. Charleston has 4 times as many sheep as Seattle. How many sheep do Toulouse, Charleston, and Seattle have together if Seattle has 20 sheep?"],
+        ["A football team played 22 games. They won 8 more than they lost. How many did they win?"],
+        ["Jim spends 2 hours watching TV and then decides to go to bed and reads for half as long. He does this 3 times a week. How many hours does he spend on TV and reading in 4 weeks?"],
+        ["Darrell and Allen's ages are in the ratio of 7:11. If their total age now is 162, calculate Allen's age 10 years from now."],
+        ["In a neighborhood, the number of rabbits pets is twelve less than the combined number of pet dogs and cats. If there are two cats for every dog, and the number of dogs is 60, how many pets in total are in the neighborhood?"],   
     ],
     cache_examples=False,
     chatbot=gr.Chatbot(