Update app.py

app.py

@@ -22,43 +22,50 @@ def gradio_predict(input_text):
         input_ids = tokenized_input["input_ids"].astype(np.int64)
         attention_mask = tokenized_input["attention_mask"].astype(np.int64)
         
-        # Create proper decoder_input_ids for autoregressive generation
+        # Initialize decoder input
         decoder_input_ids = np.array([[tokenizer.bos_token_id]], dtype=np.int64)
         
-        generated_ids = []
-        max_length = 128  # Maximum length of translation
+        print("Initial shapes:")
+        print(f"input_ids shape: {input_ids.shape}")
+        print(f"attention_mask shape: {attention_mask.shape}")
+        print(f"decoder_input_ids shape: {decoder_input_ids.shape}")
         
-        # Autoregressive generation
-        for _ in range(max_length):
-            outputs = session.run(
-                None,
-                {
-                    "input_ids": input_ids,
-                    "attention_mask": attention_mask,
-                    "decoder_input_ids": decoder_input_ids
-                }
-            )
-            
-            # Get the next token prediction
-            next_token_logits = outputs[0][0, -1, :]
-            next_token = np.argmax(next_token_logits)
-            
-            # Stop if we hit the EOS token
-            if next_token == tokenizer.eos_token_id:
-                break
-                
-            # Append the predicted token
-            generated_ids.append(next_token)
+        # Run initial inference
+        outputs = session.run(
+            None,
+            {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "decoder_input_ids": decoder_input_ids
+            }
+        )
+        
+        print("Output information:")
+        print(f"outputs type: {type(outputs)}")
+        print(f"outputs length: {len(outputs)}")
+        print(f"outputs[0] shape: {outputs[0].shape}")
+        
+        # Get logits and convert to token ids
+        logits = outputs[0]
+        token_ids = np.argmax(logits[0], axis=-1)
+        
+        # Find end of sequence
+        eos_token_id = tokenizer.eos_token_id
+        end_idx = np.where(token_ids == eos_token_id)[0]
+        if len(end_idx) > 0:
+            token_ids = token_ids[:end_idx[0]]
             
-            # Update decoder_input_ids for next iteration
-            decoder_input_ids = np.array([[tokenizer.bos_token_id] + generated_ids], dtype=np.int64)
+        print(f"token_ids shape: {token_ids.shape}")
+        print(f"token_ids: {token_ids}")
         
-        # Decode the generated sequence
-        translated_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
+        # Decode the sequence
+        translated_text = tokenizer.decode(token_ids, skip_special_tokens=True)
         return translated_text
         
     except Exception as e:
         print(f"Detailed error: {str(e)}")
+        import traceback
+        print(traceback.format_exc())
         return f"Error during translation: {str(e)}"
 
 # Gradio interface for the web app