Update app.py

app.py

@@ -1,10 +1,27 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+import threading
+import torch
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+# Load model directly
+model = AutoModelForCausalLM.from_pretrained("Oysiyl/gemma-3-1B-GRPO")
+tokenizer = AutoTokenizer.from_pretrained("Oysiyl/gemma-3-1B-GRPO")
+
+
+def process_history(history):
+    """Process chat history into the format expected by the model."""
+    processed_history = []
+    for user_msg, assistant_msg in history:
+        if user_msg:
+            processed_history.append({"role": "user", "content": [{"type": "text", "text": user_msg}]})
+        if assistant_msg:
+            processed_history.append({"role": "assistant", "content": [{"type": "text", "text": assistant_msg}]})
+    return processed_history
+
+
+def process_new_user_message(message):
+    """Process a new user message into the format expected by the model."""
+    return [{"type": "text", "text": message}]
 
 
 def respond(
@@ -15,29 +32,100 @@ def respond(
     temperature,
     top_p,
 ):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
+    # Format messages according to Gemma's expected chat format
+    messages = []
+    if system_message:
+        messages.append({"role": "system", "content": [{"type": "text", "text": system_message}]})
+    
+    messages.extend(process_history(history))
+    messages.append({"role": "user", "content": process_new_user_message(message)})
+    
+    # Apply chat template
+    inputs = tokenizer.apply_chat_template(
         messages,
-        max_tokens=max_tokens,
-        stream=True,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt"
+    )
+    
+    if torch.cuda.is_available():
+        inputs = inputs.to(device=model.device, dtype=torch.bfloat16)
+        model.to("cuda")
+    
+    # Set up the streamer
+    streamer = TextIteratorStreamer(tokenizer, timeout=30.0, skip_prompt=True, skip_special_tokens=False)
+    
+    # Run generation in a separate thread
+    generate_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
+        top_k=64,  # Recommended Gemma-3 setting
+        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"
+    for token in streamer:
+        output += token
+        # Check for various solution patterns
+        if "<SOLUTION>" in output:
+            # Original solution pattern
+            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
+            else:
+                # Handle case where closing tag is missing
+                formatted_output = (
+                    output[:solution_start] + 
+                    "Final answer: **" + output[solution_start:] + "**"
+                )
+                yield formatted_output
+        # Check if end_working_out tag is present
+        elif "</end_working_out>" in output:
+            solution_start = output.find("</end_working_out>") + len("</end_working_out>")
+            formatted_output = (
+                output[:solution_start] + 
+                "\nFinal answer: **" + output[solution_start:] + "**"
+            )
+            yield formatted_output
+        # Check if start_working_out is present but end_working_out is missing
+        elif "<start_working_out>" in output:
+            # Check if there's a SOLUTION tag after start_working_out
+            working_start = output.find("<start_working_out>")
+            if "<SOLUTION>" in output[working_start:]:
+                solution_start = output.find("<SOLUTION>", working_start) + len("<SOLUTION>")
+                solution_end = output.find("</SOLUTION>", solution_start)
+                if solution_end > solution_start:
+                    formatted_output = (
+                        output[:solution_start] + 
+                        "Final answer: **" + output[solution_start:solution_end] + "**" + 
+                        output[solution_end:]
+                    )
+                    yield formatted_output
+                else:
+                    formatted_output = (
+                        output[:solution_start] + 
+                        "Final answer: **" + output[solution_start:] + "**"
+                    )
+                    yield formatted_output
+            else:
+                # No clear solution identified
+                yield output
+        else:
+            yield output
 
 
 """
@@ -46,9 +134,12 @@ For information on how to customize the ChatInterface, peruse the gradio docs: h
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Textbox(
+            value="You are given a problem.\nThink about the problem and provide your working out.\nPlace it between <start_working_out> and <end_working_out>.\nThen, provide your solution between <SOLUTION></SOLUTION>",
+            label="System message"
+        ),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=1.0, step=0.1, label="Temperature"),
         gr.Slider(
             minimum=0.1,
             maximum=1.0,
@@ -57,8 +148,20 @@ demo = gr.ChatInterface(
             label="Top-p (nucleus sampling)",
         ),
     ],
+    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?"],
+    ],
+    chatbot=gr.Chatbot(
+        latex_delimiters=[
+            {"left": "$$", "right": "$$", "display": True},
+            {"left": "$", "right": "$", "display": False}
+        ],
+    ),
 )
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()