Update app.py

app.py

@@ -76,22 +76,34 @@ LANGUAGE_MODELS = {
     'Zulu': 'zu',
 }
 
+
 @st.cache_resource
 def load_model():
     tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")
     model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")
     return tokenizer, model
 
-def translate(text, target_language):
+def translate(text, target_language, max_chunk_size=500):
     tokenizer, model = load_model()
     
     # Set the target language code for translation
     target_lang_code = LANGUAGE_MODELS.get(target_language)
 
     tokenizer.src_lang = "en"
-    encoded_input = tokenizer(text, return_tensors="pt")
-    generated_tokens = model.generate(**encoded_input, forced_bos_token_id=tokenizer.get_lang_id(target_lang_code))
-    translation = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
+    
+    # Split the text into chunks if it's too long
+    tokens = tokenizer.encode(text, return_tensors="pt")
+    input_ids = tokens[0]
+    translations = []
+
+    for i in range(0, len(input_ids), max_chunk_size):
+        chunk = input_ids[i:i + max_chunk_size].unsqueeze(0)
+        generated_tokens = model.generate(**{'input_ids': chunk, 'forced_bos_token_id': tokenizer.get_lang_id(target_lang_code)})
+        chunk_translation = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
+        translations.append(chunk_translation)
+    
+    # Combine all translated chunks
+    translation = " ".join(translations)
     return translation
 
 st.title('Welcome to PolyTranslate')