Update app.py

app.py

@@ -1,3 +1,5 @@
+# code v13
+
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
@@ -183,14 +185,16 @@ class NLLBTranslator:
                 max_length=512
             ).to(self.device)
             
-            # Generate without forced language token first
+            # Generate without forced language token to avoid tokenizer issues
             with torch.no_grad():
                 outputs = self.model.generate(
                     **inputs,
                     max_length=512,
                     num_beams=5,
                     early_stopping=True,
-                    do_sample=False
+                    do_sample=False,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id
                 )
             
             # Decode
@@ -231,25 +235,38 @@ class NLLBTranslator:
                     max_length=512
                 ).to(self.device)
                 
-                # Get target language token ID
+                # Get target language token ID using different methods
+                target_token_id = None
                 try:
-                    target_token_id = self.tokenizer.lang_code_to_id[target_code]
-                except KeyError:
-                    logger.warning(f"Language code {target_code} not found in tokenizer, using default")
-                    target_token_id = self.tokenizer.pad_token_id
+                    # Method 1: Try lang_code_to_id
+                    if hasattr(self.tokenizer, 'lang_code_to_id'):
+                        target_token_id = self.tokenizer.lang_code_to_id[target_code]
+                    # Method 2: Try convert_tokens_to_ids
+                    elif hasattr(self.tokenizer, 'convert_tokens_to_ids'):
+                        target_token_id = self.tokenizer.convert_tokens_to_ids(target_code)
+                    # Method 3: Try getting from vocabulary
+                    else:
+                        target_token_id = self.tokenizer.get_vocab().get(target_code)
+                except (KeyError, AttributeError):
+                    logger.warning(f"Could not find target language token for {target_code}")
+                    target_token_id = None
                 
                 # Generate translation
+                generation_kwargs = {
+                    "max_length": 512,
+                    "num_beams": 4,
+                    "early_stopping": True,
+                    "do_sample": False,
+                    "pad_token_id": self.tokenizer.pad_token_id,
+                    "eos_token_id": self.tokenizer.eos_token_id
+                }
+                
+                # Only add forced_bos_token_id if we found a valid target token
+                if target_token_id is not None:
+                    generation_kwargs["forced_bos_token_id"] = target_token_id
+                
                 with torch.no_grad():
-                    translated_tokens = self.model.generate(
-                        **inputs,
-                        forced_bos_token_id=target_token_id,
-                        max_length=512,
-                        num_beams=4,
-                        early_stopping=True,
-                        do_sample=False,
-                        pad_token_id=self.tokenizer.pad_token_id,
-                        eos_token_id=self.tokenizer.eos_token_id
-                    )
+                    translated_tokens = self.model.generate(**inputs, **generation_kwargs)
                 
                 # Decode translations
                 translations = self.tokenizer.batch_decode(
@@ -257,10 +274,9 @@ class NLLBTranslator:
                     skip_special_tokens=True
                 )
                 
-                # Clean up translations (remove source language tokens if present)
+                # Clean up translations
                 cleaned_translations = []
                 for trans in translations:
-                    # Remove any language tokens that might be in the output
                     cleaned = trans.strip()
                     if cleaned:
                         cleaned_translations.append(cleaned)
@@ -290,25 +306,17 @@ class NLLBTranslator:
                             max_length=512
                         ).to(self.device)
                         
-                        # Try different approaches for target language
-                        generation_kwargs = {
-                            "max_length": 512,
-                            "num_beams": 2,
-                            "early_stopping": True,
-                            "do_sample": False,
-                            "pad_token_id": self.tokenizer.pad_token_id,
-                            "eos_token_id": self.tokenizer.eos_token_id
-                        }
-                        
-                        # Try with target language token first
-                        try:
-                            target_token_id = self.tokenizer.lang_code_to_id[target_code]
-                            generation_kwargs["forced_bos_token_id"] = target_token_id
-                        except KeyError:
-                            logger.warning(f"Target language {target_code} not in tokenizer, trying without forced_bos_token_id")
-                        
+                        # Use simple generation without forced language tokens
                         with torch.no_grad():
-                            translated_tokens = self.model.generate(**inputs, **generation_kwargs)
+                            translated_tokens = self.model.generate(
+                                **inputs,
+                                max_length=512,
+                                num_beams=2,
+                                early_stopping=True,
+                                do_sample=False,
+                                pad_token_id=self.tokenizer.pad_token_id,
+                                eos_token_id=self.tokenizer.eos_token_id
+                            )
                         
                         translation = self.tokenizer.decode(
                             translated_tokens[0],