Update app.py

app.py

@@ -78,49 +78,41 @@ def enhance_for_gradio(prompt_text, api_key):
     ]
     
     try:
-        result = pipe(messages, max_new_tokens=256, temperature=0.7, do_sample=True, pad_token_id=pipe.tokenizer.eos_token_id)
+        result = pipe(messages, max_new_tokens=256, temperature=0.7, do_sample=True)
         full_response = result[0]["generated_text"]
-        
-        # Extract only the assistant's response
-        if isinstance(full_response, str):
-            # 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(' "\'')
+
+        # Eğer full_response bir liste ise (örneğin [{'role': '...', 'content': '...'}, ...])
+        if isinstance(full_response, list):
+            # Assistant rolündeki son mesajı bul
+            assistant_messages = [msg["content"] for msg in full_response if msg["role"] == "assistant"]
+            if assistant_messages:
+                enhanced_prompt = assistant_messages[-1]
+            else:
+                enhanced_prompt = "No assistant response found."
+        else:
+            # Eğer string ise, assistant kısmını elle ayırmaya çalış
+            if '"role": "assistant"' in full_response:
+                parts = full_response.split('"role": "assistant"')
+                last_part = parts[-1]
+                if '"content":' in last_part:
+                    start = last_part.find('"content":') + len('"content":')
+                    content = last_part[start:].strip()
+                    if content.startswith('"'):
+                        end_quote = content[1:].find('"') + 1
+                        enhanced_prompt = content[1:end_quote]
                     else:
-                        enhanced_prompt = "Unable to extract enhanced prompt"
+                        enhanced_prompt = content.split("}")[0].strip(' "\'')
                 else:
-                    enhanced_prompt = "No assistant response found"
+                    enhanced_prompt = "Could not extract assistant content."
             else:
-                # 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 = "Unable to parse response"
-        else:
-            enhanced_prompt = str(full_response)
-        
-        return enhanced_prompt
+                # Fallback: sadece assistant kısmını tahmin et
+                enhanced_prompt = full_response.strip()
+
+        # Özel token'ları temizle
+        enhanced_prompt = enhanced_prompt.replace("<end_of_turn>", "").strip()
+
+        return enhanced_prompt if enhanced_prompt else "No enhanced prompt generated."
+
     except Exception as e:
         return f"Enhancement failed: {str(e)}"