Clean up code output.

app.py

@@ -27,6 +27,13 @@ except Exception as e:
     veri_tokenizer = None
 
 @spaces.GPU(duration=60)
+def truncate_at_code_end(text):
+    """Truncate text at 'CODE END' to remove repetitive content"""
+    if "CODE END" in text:
+        end_index = text.find("CODE END") + len("CODE END")
+        return text[:end_index].strip()
+    return text.strip()
+    
 def generate_response(user_message, history):
     if not veri_model or not veri_tokenizer:
         return history + [["Error", "Model not loaded properly"]]
@@ -73,6 +80,10 @@ def generate_response(user_message, history):
     
     response = veri_tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
 
+    # Truncate at CODE END to remove repetitive content
+    response = truncate_at_code_end(response)
+    
+
     if torch.cuda.is_available():
         torch.cuda.empty_cache()