Update app.py

app.py

@@ -78,24 +78,45 @@ def enhance_for_gradio(prompt_text, api_key):
     ]
     
     try:
-        result = pipe(messages, max_new_tokens=256, temperature=0.7, do_sample=True)
+        result = pipe(messages, max_new_tokens=256, temperature=0.7, do_sample=True, pad_token_id=pipe.tokenizer.eos_token_id)
         full_response = result[0]["generated_text"]
         
-        # Extract only the new generated content after the original messages
+        # Extract only the assistant's response
         if isinstance(full_response, str):
-            # Find the assistant's response after the user message
-            parts = full_response.split("<end_of_turn>")
-            if len(parts) > 2:
-                enhanced_prompt = parts[-1].strip()
+            # Split by the conversation structure and get the last assistant response
+            if "assistant" in full_response:
+                # Find the last occurrence of assistant role
+                assistant_parts = full_response.split('"role": "assistant"')
+                if len(assistant_parts) > 1:
+                    last_part = assistant_parts[-1]
+                    # Extract content between quotes
+                    if '"content":' in last_part:
+                        content_start = last_part.find('"content":') + len('"content":')
+                        content_part = last_part[content_start:].strip()
+                        # Find the actual content between quotes
+                        if content_part.startswith('"') or content_part.startswith("'"):
+                            quote_char = content_part[0]
+                            end_quote = content_part.find(quote_char, 1)
+                            if end_quote > 0:
+                                enhanced_prompt = content_part[1:end_quote]
+                            else:
+                                enhanced_prompt = content_part[1:].split('}')[0].strip()
+                        else:
+                            enhanced_prompt = content_part.split('}')[0].strip(' "\'')
+                    else:
+                        enhanced_prompt = "Unable to extract enhanced prompt"
+                else:
+                    enhanced_prompt = "No assistant response found"
             else:
-                # Fallback: try to extract after last user message
-                user_msg_end = full_response.rfind(prompt_text)
-                if user_msg_end != -1:
-                    enhanced_prompt = full_response[user_msg_end + len(prompt_text):].strip()
-                    # Clean up any remaining tokens
-                    enhanced_prompt = enhanced_prompt.replace("<end_of_turn>", "").strip()
+                # Fallback: try to get content after user input
+                user_content = prompt_text
+                if user_content in full_response:
+                    response_start = full_response.rfind(user_content) + len(user_content)
+                    remaining = full_response[response_start:].strip()
+                    # Clean up response
+                    enhanced_prompt = remaining.replace('"', '').replace("'", "").replace("}", "").strip()[:500]
                 else:
-                    enhanced_prompt = full_response
+                    enhanced_prompt = "Unable to parse response"
         else:
             enhanced_prompt = str(full_response)