Update cultural_model.py

cultural_model.py

@@ -1,24 +1,23 @@
 from transformers import M2M100ForConditionalGeneration
 
 class CulturalM2M100(M2M100ForConditionalGeneration):
-    """
-    Custom M2M100 model with cultural preservation features
-    Inherits from M2M100ForConditionalGeneration and adds:
-    - Cultural context tokens
-    - Special generation parameters
-    """
+    """Custom model with cultural preservation features"""
     
     def __init__(self, config):
         super().__init__(config)
-        self.cultural_token_id = 250001  # Should match tokenizer
-        
+        self.cultural_token_id = 250001  # Must match tokenizer
+    
     def generate(self, *args, **kwargs):
-        """Override generate method with cultural preservation"""
-        if kwargs.pop("cultural_preservation", False):
+        """Override generation with cultural context"""
+        cultural_preservation = kwargs.pop("cultural_preservation", False)
+        
+        if cultural_preservation:
+            # Force cultural context token
             kwargs["forced_bos_token_id"] = self.get_cultural_token()
-            kwargs["max_length"] = kwargs.get("max_length", 200)  # Longer for context
+            kwargs["max_length"] = kwargs.get("max_length", 512) + 50
+        
         return super().generate(*args, **kwargs)
     
     def get_cultural_token(self):
-        """Get the special cultural context token ID"""
+        """Get cultural preservation token ID"""
         return self.cultural_token_id