Update app.py

app.py

@@ -22,15 +22,10 @@ def gradio_predict(input_text):
         input_ids = tokenized_input["input_ids"].astype(np.int64)
         attention_mask = tokenized_input["attention_mask"].astype(np.int64)
         
-        # Initialize decoder_input_ids with start token
+        # Initialize decoder_input_ids
         decoder_input_ids = np.zeros((1, 512), dtype=np.int64)
         decoder_input_ids[:, 0] = tokenizer.bos_token_id or tokenizer.pad_token_id
         
-        print("Input values:")
-        print(f"First few input_ids: {input_ids[0][:10]}")
-        print(f"First few attention_mask: {attention_mask[0][:10]}")
-        print(f"First few decoder_input_ids: {decoder_input_ids[0][:10]}")
-        
         # Run inference
         outputs = session.run(
             None,
@@ -41,27 +36,22 @@ def gradio_predict(input_text):
             }
         )
         
-        print("Output shape and type:")
-        print(f"Output type: {type(outputs)}")
-        print(f"Output[0] type: {type(outputs[0])}")
-        print(f"Output[0] shape: {outputs[0].shape}")
+        # Process logits to get token ids
+        logits = outputs[0]  # Shape: (1, 512, vocab_size)
+        token_ids = np.argmax(logits, axis=-1)[0]  # Get token ids for first sequence
         
-        # Process outputs more carefully
-        output_ids = outputs[0]
-        if isinstance(output_ids, np.ndarray):
-            output_ids = output_ids[0]  # Take first sequence
-            # Convert to list of integers if needed
-            if isinstance(output_ids, np.ndarray):
-                output_ids = output_ids.tolist()
+        # Find where the sequence ends (pad token or eos token)
+        eos_token_id = tokenizer.eos_token_id or tokenizer.pad_token_id
+        end_idx = np.where(token_ids == eos_token_id)[0]
+        if len(end_idx) > 0:
+            token_ids = token_ids[:end_idx[0]]
         
         # Decode output
-        translated_text = tokenizer.decode(output_ids, skip_special_tokens=True)
+        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