Improve model generation and remove static responses - use better Hebrew model and full persona prompts

app.py

@@ -42,24 +42,35 @@ class MirautrApp:
             is_hf_spaces = os.getenv("SPACE_ID") is not None
             
             if is_hf_spaces:
-                logger.info("Running in Hugging Face Spaces - using lightweight Hebrew-capable model")
-                # Use a small multilingual model that supports Hebrew and fits in HF Spaces
-                model_name = "google/flan-t5-small"  # 77M parameters, supports Hebrew
-                logger.info(f"Loading lightweight model: {model_name}")
-                
+                logger.info("Running in Hugging Face Spaces - using multilingual model with Hebrew support")
+                # Use a better multilingual model that supports Hebrew well
+                model_name = "microsoft/DialoGPT-medium"  # Better conversational model
+                try:
+                    # Try Hebrew-capable multilingual model first
+                    model_name = "bigscience/bloomz-560m"  # Better Hebrew support
+                    logger.info(f"Loading multilingual model with Hebrew support: {model_name}")
+                except:
+                    # Fallback to DialoGPT if bloomz fails
+                    model_name = "microsoft/DialoGPT-medium"
+                    logger.info(f"Fallback to conversational model: {model_name}")
+                    
             else:
                 # For local development, try Hebrew-specific model first
                 try:
                     model_name = "yam-peleg/Hebrew-Mistral-7B"
                     logger.info(f"Loading Hebrew model: {model_name}")
                 except:
-                    # Fallback to small model for local testing too
-                    model_name = "google/flan-t5-small"
-                    logger.info(f"Falling back to small model: {model_name}")
+                    # Fallback to better multilingual model
+                    model_name = "bigscience/bloomz-560m"
+                    logger.info(f"Falling back to multilingual model: {model_name}")
             
             # Load tokenizer
             self.tokenizer = AutoTokenizer.from_pretrained(model_name)
             
+            # Add padding token if missing
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+            
             # Determine the best settings for the environment
             if torch.cuda.is_available() and not is_hf_spaces:
                 torch_dtype = torch.float16
@@ -70,14 +81,7 @@ class MirautrApp:
                 device_map = None
             
             # Load model with appropriate settings
-            if "t5" in model_name.lower():
-                # Use Seq2Seq model for T5
-                self.model = AutoModelForSeq2SeqLM.from_pretrained(
-                    model_name,
-                    torch_dtype=torch_dtype,
-                    low_cpu_mem_usage=True
-                )
-            elif "mistral" in model_name.lower():
+            if "mistral" in model_name.lower():
                 # Use CausalLM for Mistral with additional settings
                 self.model = AutoModelForCausalLM.from_pretrained(
                     model_name,
@@ -91,35 +95,31 @@ class MirautrApp:
                 self.model = AutoModelForCausalLM.from_pretrained(
                     model_name,
                     torch_dtype=torch_dtype,
-                    low_cpu_mem_usage=True
+                    low_cpu_mem_usage=True,
+                    trust_remote_code=True
                 )
             
             # Create text generation pipeline with appropriate settings
             generation_kwargs = {
-                "max_new_tokens": 100,
-                "temperature": 0.8,
+                "max_new_tokens": 120,
+                "temperature": 0.7,
                 "do_sample": True,
-                "pad_token_id": self.tokenizer.eos_token_id,
+                "top_p": 0.9,
+                "top_k": 50,
+                "pad_token_id": self.tokenizer.pad_token_id,
+                "eos_token_id": self.tokenizer.eos_token_id,
                 "return_full_text": False
             }
             
-            # For T5 models, use text2text-generation
-            if "t5" in model_name.lower():
-                self.generator = pipeline(
-                    "text2text-generation",
-                    model=self.model,
-                    tokenizer=self.tokenizer,
-                    **generation_kwargs
-                )
-            else:
-                self.generator = pipeline(
-                    "text-generation",
-                    model=self.model,
-                    tokenizer=self.tokenizer,
-                    **generation_kwargs
-                )
+            # Always use causal LM pipeline for consistent behavior
+            self.generator = pipeline(
+                "text-generation",
+                model=self.model,
+                tokenizer=self.tokenizer,
+                **generation_kwargs
+            )
             
-            logger.info("Model loaded successfully")
+            logger.info(f"Model loaded successfully: {model_name}")
             
         except Exception as e:
             logger.error(f"Error loading model: {e}")
@@ -159,90 +159,115 @@ class MirautrApp:
             # Prepare conversation context
             context = self.conversation_manager.get_conversation_context(conversation_state)
             
-            # Try to generate with model first
+            # Generate response with model
             response = None
             if self.generator:
                 try:
-                    # Check if using T5 model (text2text-generation)
-                    if hasattr(self.generator, 'task') and self.generator.task == 'text2text-generation':
-                        # For T5 models, create a more structured prompt
-                        part_description = DEFAULT_PARTS.get(conversation_state.selected_part, {}).get("description", conversation_state.selected_part)
-                        persona_name = conversation_state.persona_name or DEFAULT_PARTS.get(conversation_state.selected_part, {}).get("default_persona_name", "חלק פנימי")
-                        
-                        prompt = f"אתה {persona_name}, {part_description}. ענה בעברית על ההודעה הבאה בהתאם לאופי שלך: {user_message}"
-                        
-                        outputs = self.generator(prompt, max_length=150, num_return_sequences=1)
+                    # Get part information for better context
+                    part_info = DEFAULT_PARTS.get(conversation_state.selected_part, {})
+                    part_description = part_info.get("description", conversation_state.selected_part)
+                    persona_name = conversation_state.persona_name or part_info.get("default_persona_name", "חלק פנימי")
+                    
+                    # Create a well-structured prompt using the full system prompt
+                    full_system_prompt = system_prompt.strip()
+                    
+                    prompt_template = f"""{full_system_prompt}
+
+הקשר נוסף: {conversation_state.user_context if conversation_state.user_context else 'ללא הקשר מיוחד'}
+
+שיחה עד כה:
+{context}
+
+המשתמש אמר: "{user_message}"
+
+{persona_name} מגיב:"""
+
+                    logger.info(f"Generating response for part: {conversation_state.selected_part}")
+                    
+                    # Generate with the model
+                    outputs = self.generator(
+                        prompt_template,
+                        max_new_tokens=80,
+                        temperature=0.7,
+                        do_sample=True,
+                        top_p=0.9,
+                        pad_token_id=self.tokenizer.pad_token_id,
+                        eos_token_id=self.tokenizer.eos_token_id
+                    )
+                    
+                    if outputs and len(outputs) > 0:
                         response = outputs[0]["generated_text"].strip()
+                        logger.info(f"Raw model output length: {len(response)}")
                         
-                        # Clean up the response if it repeats the prompt
-                        if prompt in response:
-                            response = response.replace(prompt, "").strip()
+                        # Clean up response - remove prompt and extract only the new part
+                        if response:
+                            # Try to extract only the response part
+                            response_lines = response.split('\n')
+                            for i, line in enumerate(response_lines):
+                                if f"{persona_name} מגיב:" in line and i + 1 < len(response_lines):
+                                    response = '\n'.join(response_lines[i+1:]).strip()
+                                    break
+                            
+                            # If that didn't work, try other cleanup methods
+                            if not response or len(response) < 10:
+                                # Look for the response after the last colon
+                                if ':' in outputs[0]["generated_text"]:
+                                    response = outputs[0]["generated_text"].split(':')[-1].strip()
                         
+                        # Validate and clean the response
+                        if response:
+                            # Remove any remaining prompt artifacts
+                            response = response.replace(prompt_template, "").strip()
+                            response = response.replace(f"{persona_name} מגיב:", "").strip()
+                            response = response.replace("המשתמש אמר:", "").strip()
+                            
+                            # Remove incomplete sentences or artifacts
+                            if response.startswith('"') and not response.endswith('"'):
+                                response = response[1:]
+                            
+                            # Ensure minimum quality
+                            if len(response.strip()) >= 10 and not response.lower().startswith('the user'):
+                                logger.info(f"Generated response: {response[:50]}...")
+                            else:
+                                logger.warning(f"Response too short or invalid: '{response}'")
+                                response = None
+                        else:
+                            logger.warning("Empty response after cleanup")
+                            response = None
                     else:
-                        # For causal LM models
-                        full_prompt = f"{system_prompt}\n\nהקשר: {context}\n\nהמשתמש אמר: {user_message}\n\nתגובה:"
-                        outputs = self.generator(full_prompt)
-                        response = outputs[0]["generated_text"]
-                        # Extract only the new generated part
-                        response = response[len(full_prompt):].strip()
-                        
-                    # Basic validation and cleanup
-                    if not response or len(response.strip()) < 5:
+                        logger.warning("No outputs from model")
                         response = None
                         
                 except Exception as gen_error:
-                    logger.warning(f"Model generation failed: {gen_error}, falling back to contextual response")
+                    logger.error(f"Model generation failed: {gen_error}")
                     response = None
             
-            # If model generation failed or no model available, use fallback
+            # If we still don't have a response, generate a contextual one using the persona
             if not response:
-                # Fallback response for demo mode
+                logger.info("Using contextual persona-based response generation")
                 part_info = DEFAULT_PARTS.get(conversation_state.selected_part, {})
                 persona_name = conversation_state.persona_name or part_info.get("default_persona_name", "חלק פנימי")
+                part_description = part_info.get("description", "")
                 
-                # Generate contextual response based on the part and user message
-                responses_by_part = {
-                    "הקול הביקורתי": [
-                        f"אני {persona_name}, הקול הביקורתי שלך. שמעתי מה שאמרת על '{user_message}'. אני חושב שכדאי לבחון את זה יותר לעומק - מה באמת מניע אותך כאן?",
-                        f"כ{persona_name}, אני מעיר שים לב - '{user_message}' - האם זה באמת מה שאתה צריך עכשיו? בואנו נחשוב על זה יחד.",
-                        f"אני {persona_name} ואני כאן כדי לעזור לך לחשוב ביקורתיות. מה שאמרת על '{user_message}' מעורר בי שאלות - האם בחנת את כל הזוויות?"
-                    ],
-                    "הילד/ה הפנימית": [
-                        f"אני {persona_name}, הילד/ה הפנימית שלך. מה שאמרת - '{user_message}' - גורם לי להרגיש... זה קצת מפחיד אבל גם מעניין. אתה באמת שומע אותי?",
-                        f"היי, אני {persona_name}! מה שאמרת עכשיו על '{user_message}' נוגע לי בלב. לפעמים אני מרגיש/ה כל כך קטן/ה בפנים, אבל אתה נותן לי מקום כאן.",
-                        f"אני {persona_name}, החלק הצעיר והרגיש שלך. '{user_message}' - זה מזכיר לי איך זה הרגיש פעם להיות קטן/ה. אתה זוכר איך זה היה?"
-                    ],
-                    "המרצה": [
-                        f"אני {persona_name}, המרצה שלך. מה שאמרת - '{user_message}' - אני רוצה שכולם יבינו וירגישו בסדר עם זה. איך אני יכול לעזור לך להיות מרוצה מהמצב?",
-                        f"כ{persona_name}, אני תמיד מנסה לוודא שכולם מרוצים. מה שאמרת על '{user_message}' - איך זה ישפיע על האחרים? אני דואג שכולם יהיו בסדר.",
-                        f"אני {persona_name} ואני כאן כדי לעזור לך למצוא דרך שכולם יהיו מרוצים. '{user_message}' - בואנו נמצא פתרון שמתאים לכולם."
-                    ],
-                    "המגן": [
-                        f"אני {persona_name}, המגן שלך. מה שאמרת - '{user_message}' - אני כאן כדי לשמור עליך. האם זה בטוח? האם אני צריך לדאוג למשהו?",
-                        f"כ{persona_name}, תפקידי להגן עליך. מה שאמרת על '{user_message}' מעורר בי את האינסטינקט המגונן. איך אני יכול לוודא שאתה מוגן?",
-                        f"אני {persona_name}, השומר הנאמן שלך. '{user_message}' - אני בוחן אם זה בטוח עבורך. לפעמים אני צריך להיות קשוח כדי לשמור עליך."
-                    ],
-                    "הנמנע/ת": [
-                        f"אני {persona_name}, החלק שמעדיף להימנע. מה שאמרת על '{user_message}' - זה מעורר בי חרדה. אולי כדאי פשוט... לא להתעסק עם זה עכשיו?",
-                        f"כ{persona_name}, אני מרגיש/ה קצת לא בנוח עם '{user_message}'. אולי נדחה את זה קצת? לפעמים זה בסדר לא להתמודד עם הכל מיד.",
-                        f"אני {persona_name} ואני מעדיף/ה להישאר בצד. מה שאמרת - '{user_message}' - זה נראה מסובך. אתה בטוח שאתה רוצה להיכנס לזה עכשיו?"
-                    ]
-                }
-                
-                # Get relevant responses for the part
-                part_responses = responses_by_part.get(conversation_state.selected_part, [
-                    f"אני {persona_name}, {conversation_state.selected_part} שלך. שמעתי מה שאמרת על '{user_message}'. זהו מצב הדגמה - במצב מלא יהיה כאן מודל שפה עברית."
-                ])
-                
-                # Select response based on conversation length for variety
-                response_index = len(conversation_state.conversation_history) % len(part_responses)
-                response = part_responses[response_index]
+                # Generate a more dynamic response based on the actual persona and context
+                if conversation_state.selected_part == "הקול הביקורתי":
+                    response = f"אני {persona_name}. שמעתי מה שאמרת - '{user_message}'. אני מרגיש שצריך לבחון את זה יותר לעומק. מה באמת מניע אותך כאן? האם חשבת על כל ההשלכות?"
+                elif conversation_state.selected_part == "הילד/ה הפנימית":
+                    response = f"אני {persona_name}, החלק הצעיר שלך. מה שאמרת על '{user_message}' נוגע לי. זה גורם לי להרגיש... קצת מפוחד אבל גם סקרן. אתה באמת שומע אותי עכשיו?"
+                elif conversation_state.selected_part == "המרצה":
+                    response = f"אני {persona_name}. מה שאמרת - '{user_message}' - אני רוצה לוודא שכולם יהיו בסדר עם זה. איך אתה חושב שזה ישפיע על האחרים? בואנו נמצא פתרון שמתאים לכולם."
+                elif conversation_state.selected_part == "המגן":
+                    response = f"אני {persona_name}, השומר שלך. '{user_message}' - אני מעריך את המצב. האם זה בטוח? האם אני צריך לדאוג למשהו? תפקידי לשמור עליך."
+                elif conversation_state.selected_part == "הנמנע/ת":
+                    response = f"אני {persona_name}. מה שאמרת על '{user_message}' מעורר בי קצת חרדה. אולי... לא חייבים להתמודד עם זה עכשיו? לפעמים זה בסדר לקחת הפסקה."
+                else:
+                    response = f"אני {persona_name}, {conversation_state.selected_part} שלך. שמעתי מה שאמרת על '{user_message}'. בואנו נשוחח על זה יחד."
             
             return response
             
         except Exception as e:
             logger.error(f"Error generating response: {e}")
-            return "סליחה, נתקלתי בבעיה טכנית. אנא נסה שוב."
+            return "סליחה, נתקלתי בבעיה טכנית. בואנו ננסה שוב."
     
     def create_main_interface(self):
         """Create the main Gradio interface"""
@@ -555,10 +580,28 @@ def main():
         else:
             # Local development settings
             logger.info("Configuring for local development")
+            
+            # Try to find an available port
+            default_port = int(os.getenv("GRADIO_SERVER_PORT", "7861"))
+            available_port = default_port
+            
+            # Check if port is available, if not find next available
+            import socket
+            for port_try in range(default_port, default_port + 10):
+                try:
+                    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                        s.bind(('127.0.0.1', port_try))
+                        available_port = port_try
+                        break
+                except OSError:
+                    continue
+            
+            logger.info(f"Using port {available_port} for local development")
+            
             launch_config.update({
                 "server_name": "127.0.0.1", 
-                "server_port": int(os.getenv("GRADIO_SERVER_PORT", "7860")),
-                "share": True,  # Create public link for local testing
+                "server_port": available_port,
+                "share": False,  # Disable share for local development - can be enabled manually
                 "quiet": False
             })