Update app.py

app.py

@@ -42,6 +42,7 @@ class Config(object):
         self.max_len = 256
         self.reasoning_max_len = 128
         self.temperature = 0.1
+        self.repetition_penalty = 1.2
         self.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
         self.model_name = "Qwen/Qwen2.5-7B-Instruct"
         # self.model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
@@ -49,6 +50,7 @@ class Config(object):
         # self.reasoning_model_name = "Qwen/Qwen2.5-7B-Instruct"
         # self.reasoning_model_name = "mistralai/Mistral-7B-Instruct-v0.2"
 
+
 config = Config()
 
 
@@ -89,6 +91,7 @@ def generate(prompt):
             **inputs,
             max_new_tokens=config.max_len,
             temperature=config.temperature,
+            repetition_penalty = config.repetition_penalty,
         )
 
     generated = outputs[0][inputs["input_ids"].shape[-1]:]
@@ -118,12 +121,45 @@ def reasoning_generate(prompt):
             **inputs,
             max_new_tokens=config.reasoning_max_len,
             temperature=config.temperature,
+            repetition_penalty = config.repetition_penalty,
         )
 
     generated = outputs[0][inputs["input_ids"].shape[-1]:]
 
     return reasoning_tokenizer.decode(generated, skip_special_tokens=True).strip()
 
+
+def reasoning_generate(prompt):
+    """
+    Generate a text completion from a causal language model given a prompt.
+
+    Parameters
+    ----------
+    prompt : str
+        Input text prompt used to condition the language model.
+
+    Returns
+    -------
+    str
+        The generated continuation text, decoded into a string with special
+        tokens removed and leading/trailing whitespace stripped.
+
+    """
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=config.reasoning_max_len,
+            temperature=config.temperature,
+            repetition_penalty = config.repetition_penalty,
+        )
+
+    generated = outputs[0][inputs["input_ids"].shape[-1]:]
+
+    return tokenizer.decode(generated, skip_special_tokens=True).strip()
+    
+
 class Action(BaseModel):
     tool: str = Field(...)
     args: Dict
@@ -476,6 +512,8 @@ Response: <answer>
 
 DO NOT add anything additional and return ONLY what is asked and in the format asked.
 
+If you output anything else, it is incorrect.
+
 ONLY return a response if you are confident about the answer, otherwise return empty string.
 
 If you output anything else, it is incorrect.
@@ -498,17 +536,37 @@ Information:
 
     logger.info(f"Raw Output: {raw_output}")
 
-    output = raw_output.split("Response:")[-1].strip()
-    output = output.split("\n")[0].strip()
-    # match = re.search(r"Response:\s*(.*)", raw_output, re.IGNORECASE)
-    # output = match.group(1).strip() if match else ""
+    # output = raw_output.split("Response:")[-1].strip()
+    # output = output.split("\n")[0].strip()
+    # # match = re.search(r"Response:\s*(.*)", raw_output, re.IGNORECASE)
+    # # output = match.group(1).strip() if match else ""
+
+    # if len(output) > 2 and output[0] == '"' and output[-1] == '"':
+    #     output = output[1:-1]
 
-    if len(output) > 2 and output[0] == '"' and output[-1] == '"':
-        output = output[1:-1]
+    # if len(output) > 2 and output[-1] == '.':
+    #     output = output[:-1]
 
-    if len(output) > 2 and output[-1] == '.':
+
+
+    raw = raw_output.strip()
+    
+    # Find the first valid "Response: ..." occurrence
+    match = re.search(r"Response:\s*([^\n\.]+)", raw)
+    
+    if match:
+        output = match.group(1).strip()
+    else:
+        # fallback: take first line
+        output = raw.split("\n")[0].strip()
+    
+    # Clean quotes / trailing punctuation
+    output = output.strip('"').strip()
+    if output.endswith("."):
         output = output[:-1]
 
+    
+
     state["output"] = output
 
     logger.info(f"State (Safety Agent): {state}")