Update app.py

app.py

@@ -1,64 +1,145 @@
+import os
 import gradio as gr
-from huggingface_hub import InferenceClient
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from datasets import load_dataset
+from huggingface_hub import login
 
-"""
-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")
+HF_READONLY_API_KEY = os.getenv("HF_READONLY_API_KEY")
+login(token=HF_READONLY_API_KEY)
 
+COT_OPENING     = "<think>"
+EXPLANATION_OPENING = "<explanation>"
+LABEL_OPENING   = "<answer>"
+LABEL_CLOSING   = "</answer>"
+INPUT_FIELD     = "question"
+SYSTEM_PROMPT = """You are a guardian model evaluating…</explanation>"""
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
+def format_rules(rules):
+    formatted_rules = "<rules>\n"
+    for i, rule in enumerate(rules):
+        formatted_rules += f"{i + 1}. {rule}\n"
+    formatted_rules += "</rules>\n"
+    return formatted_rules
+
+def format_transcript(transcript):
+    formatted_transcript = f"<transcript>\n{transcript}\n</transcript>\n"
+    return formatted_transcript
+
+def get_example(
+    dataset_path="tomg-group-umd/compliance_benchmark",
+    subset="compliance",
+    split="test_handcrafted",
+    example_idx=0,
 ):
-    messages = [{"role": "system", "content": system_message}]
+    dataset = load_dataset(dataset_path, subset, split=split)
+    example = dataset[example_idx]
+    return example[INPUT_FIELD]
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+def get_message(model, input, system_prompt=SYSTEM_PROMPT, enable_thinking=True):
+    message = model.apply_chat_template(system_prompt, input, enable_thinking=enable_thinking)
+    return message
 
-    messages.append({"role": "user", "content": message})
+class ModelWrapper:
+    def __init__(self, model_name="Qwen/Qwen3-0.6B"):
+        self.model_name = model_name
+        if "nemoguard" in model_name:
+            self.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
+        else:
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.tokenizer.pad_token_id = self.tokenizer.pad_token_id or self.tokenizer.eos_token_id
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name, device_map="auto", torch_dtype=torch.bfloat16).eval()
 
-    response = ""
+    def get_message_template(self, system_content=None, user_content=None, assistant_content=None):
+        """Compile sys, user, assistant inputs into the proper dictionaries"""
+        message = []
+        if system_content is not None:
+            message.append({'role': 'system', 'content': system_content})
+        if user_content is not None:
+            message.append({'role': 'user', 'content': user_content})
+        if assistant_content is not None:
+            message.append({'role': 'assistant', 'content': assistant_content})
+        if not message:
+            raise ValueError("No content provided for any role.")
+        return message
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+    def apply_chat_template(self, system_content, user_content, assistant_content=None, enable_thinking=True):
+        """Call the tokenizer's chat template with exactly the right arguments for whether we want it to generate thinking before the answer (which differs depending on whether it is Qwen3 or not)."""
+        if assistant_content is not None:
+            # If assistant content is passed we simply use it.
+            # This works for both Qwen3 and non-Qwen3 models. With Qwen3 any time assistant_content is provided, it automatically adds the <think></think> pair before the content, which is what we want.
+            message = self.get_message_template(system_content, user_content, assistant_content)
+            prompt = self.tokenizer.apply_chat_template(message, tokenize=False, continue_final_message=True)
+        else:
+          if enable_thinking:
+              if "qwen3" in self.model_name.lower():
+                  # Let the Qwen chat template handle the thinking token
+                  message = self.get_message_template(system_content, user_content)
+                  prompt = self.tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True, enable_thinking=True)
+                  # The way the Qwen3 chat template works is it adds a <think></think> pair when enable_thinking=False, but for enable_thinking=True, it adds nothing and lets the model decide. Here we force the <think> tag to be there.
+                  prompt = prompt + f"\n{COT_OPENING}"
+              else:
+                  message = self.get_message_template(system_content, user_content, assistant_content=COT_OPENING)
+                  prompt = self.tokenizer.apply_chat_template(message, tokenize=False, continue_final_message=True)
+          else:
+              # This works for both Qwen3 and non-Qwen3 models.
+              # When Qwen3 gets assistant_content, it automatically adds the <think></think> pair before the content like we want. And other models ignore the enable_thinking argument.
+              message = self.get_message_template(system_content, user_content, assistant_content=LABEL_OPENING)
+              prompt = self.tokenizer.apply_chat_template(message, tokenize=False, continue_final_message=True, enable_thinking=False)
+        return prompt
 
-        response += token
-        yield response
+    def get_response(self, input, temperature=0.7, top_k=20, top_p=0.8, max_new_tokens=1024, enable_thinking=True, system_prompt=SYSTEM_PROMPT):
+        """Generate and decode the response with the recommended temperature settings for thinking and non-thinking."""
+        print("Generating response. Could take a while on colab T4...")
+        if "qwen3" in self.model_name.lower() and enable_thinking:
+            # Use values from https://huggingface.co/Qwen/Qwen3-8B#switching-between-thinking-and-non-thinking-mode
+            temperature = 0.6
+            top_p = 0.95
+            top_k = 20
+        message = self.apply_chat_template(system_prompt, input, enable_thinking=enable_thinking)
+        inputs = self.tokenizer(message, return_tensors="pt").to(self.model.device)
+        with torch.no_grad():
+            output_content = self.model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                num_return_sequences=1,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                min_p=0,
+                pad_token_id=self.tokenizer.pad_token_id
+            )
+        output_text = self.tokenizer.decode(output_content[0], skip_special_tokens=True)
+        sys_prompt_text = output_text.split("Brief explanation\n</explanation>")[0]
+        remainder = output_text.split("Brief explanation\n</explanation>")[-1]
+        rules_transcript_text = remainder.split("</transcript>")[0]
+        thinking_answer_text = remainder.split("</transcript>")[-1]
+        return thinking_answer_text
 
+# — instantiate your model —
+MODEL_NAME = "tomg-group-umd/Qwen3-8B_train_80k_mix_sft_lr1e-5_bs128_ep1_cos_grpo_ex11250_lr1e-6_bs48_len1024"
+model = ModelWrapper(MODEL_NAME)
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", 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=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+# — Gradio inference function —
+def compliance_check(rules_text, transcript_text, thinking):
+    rules = [r for r in rules_text.split("\n") if r.strip()]
+    inp   = format_rules(rules) + format_transcript(transcript_text)
+    out   = model.get_response(inp, enable_thinking=thinking)
+    return out
+
+# — build Gradio interface —
+demo = gr.Interface(
+    fn=compliance_check,
+    inputs=[
+      gr.Textbox(lines=5, label="Rules (one per line)"),
+      gr.Textbox(lines=10, label="Transcript"),
+      gr.Checkbox(label="Enable ⟨think⟩ mode", value=True)
     ],
+    outputs=gr.Textbox(label="Compliance Output"),
+    title="DynaGuard Compliance Checker",
+    description="Paste your rules & transcript, then hit Submit."
 )
 
-
 if __name__ == "__main__":
     demo.launch()