model output answer extraction

utils/llama_utils.py

@@ -2,6 +2,7 @@ import os
 import torch
 from transformers import BitsAndBytesConfig, AutoModelForCausalLM, PreTrainedTokenizerFast
 from peft import PeftModel
+import re
 import streamlit as st
 
 # Set the cache directory to persistent storage
@@ -128,6 +129,9 @@ def generate_response(
         return_tensors="pt"
     ).to(model.device)
 
+    # Store the number of input tokens for reference
+    input_token_length = inputs.input_ids.shape[1]
+
     generation_params = {
         "do_sample": do_sample,
         "temperature": temperature if do_sample else None,
@@ -139,10 +143,19 @@ def generate_response(
 
     output = model.generate(**inputs, **generation_params)
 
-    # Decode and clean up response
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    # Extract only the newly generated tokens
+    new_tokens = output[0][input_token_length:]
+     
+    # Decode only the new tokens
+    response = tokenizer.decode(new_tokens, skip_special_tokens=True)
 
     if 'assistant' in response:
         response = response.split('assistant')[1].strip()
 
+    # In case there's still any assistant prefix, clean it up
+    if response.startswith("assistant") or response.startswith("<assistant>"):
+        response = re.sub(r"^assistant[:\s]*|^<assistant>[\s]*", "", response, flags=re.IGNORECASE)
+    
+    response = re.sub(r'^\s*(?:answer\s*)+:?\s*', '', response, flags=re.IGNORECASE)
+    
     return response