Update app_low.py

app_low.py

@@ -52,38 +52,53 @@ def enhance_prompt(user_prompt, temperature, max_tokens, chat_history):
 
     return chat_history
 
+import re
+
 def extract_later_part(user_prompt, generated_text):
     """
-    Remove the user prompt from the start of the generated text
-    and return only the new added details.
+    Cleans the model output and extracts only the enhanced (later) portion.
+    Removes prompt echoes and system tags like <end_of_turn>, <start_of_turn>, etc.
     """
-    # Ensure matching is case-insensitive and strip extra spaces
+    # Step 1: Clean up internal tags
+    cleaned = re.sub(r"<.*?>", "", generated_text)  # Remove <end_of_turn>, <start_of_turn>, etc.
+    cleaned = cleaned.strip()
+
+    # Step 2: Normalize spaces
+    cleaned = re.sub(r"\s+", " ", cleaned)
+
+    # Step 3: Try removing the original prompt if repeated
     user_prompt_clean = user_prompt.strip().lower()
-    gen_clean = generated_text.strip()
+    cleaned_lower = cleaned.lower()
+
+    if cleaned_lower.startswith(user_prompt_clean):
+        cleaned = cleaned[len(user_prompt):].strip(",. ").strip()
+
+    return cleaned
 
-    if gen_clean.lower().startswith(user_prompt_clean):
-        return gen_clean[len(user_prompt_clean):].strip(", ").strip()
-    return gen_clean
 
 # ===================== Prompt Enhancer Function =====================
 def enhance_prompt1(user_prompt, temperature, max_tokens, chat_history):
-    # Build the messages for the model
     messages = [
         {"role": "system", "content": "Enhance and expand the following prompt with more details and context:"},
         {"role": "user", "content": user_prompt}
     ]
-    
-    # Generate model prompt
+
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     output = pipe(prompt, max_new_tokens=256)
-    enhanced_text = output[0]['generated_text']
-    print(enhanced_text)
-    enhanced_text = output[0]['generated_text']
-    later_part = extract_later_part(user_prompt, enhanced_text)
-    chat_history.append({"role": "user", "content": user_prompt})
-    chat_history.append({"role": "assistant", "content": later_part})
+    raw_output = output[0]['generated_text']
+
+    print("=== RAW MODEL OUTPUT ===")
+    print(raw_output)
 
+    # Extract the cleaned, later portion
+    later_part = extract_later_part(user_prompt, raw_output)
+    print("=== EXTRACTED CLEANED OUTPUT ===")
+    print(later_part)
 
+    # Append to chat history for Gradio
+    chat_history = chat_history or []
+    chat_history.append({"role": "user", "content": user_prompt})
+    chat_history.append({"role": "assistant", "content": later_part})
 
     return chat_history