Update app.py

app.py

@@ -13,7 +13,7 @@ def gradio_predict(input_text):
         tokenized_input = tokenizer(
             input_text,
             return_tensors="np",
-            padding='max_length',
+            padding=True,
             truncation=True,
             max_length=512
         )
@@ -22,38 +22,57 @@ 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
-        decoder_input_ids = np.zeros((1, 512), dtype=np.int64)
-        decoder_input_ids[:, 0] = tokenizer.bos_token_id or tokenizer.pad_token_id
+        # Create proper decoder_input_ids for autoregressive generation
+        decoder_input_ids = np.array([[tokenizer.bos_token_id]], dtype=np.int64)
         
-        # Run inference
-        outputs = session.run(
-            None,
-            {
-                "input_ids": input_ids,
-                "attention_mask": attention_mask,
-                "decoder_input_ids": decoder_input_ids
-            }
-        )
-        
-        # 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
+        generated_ids = []
+        max_length = 128  # Maximum length of translation
         
-        # 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]]
+        # 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)
+            
+            # Update decoder_input_ids for next iteration
+            decoder_input_ids = np.array([[tokenizer.bos_token_id] + generated_ids], dtype=np.int64)
         
-        # Decode output
-        translated_text = tokenizer.decode(token_ids, skip_special_tokens=True)
+        # Decode the generated sequence
+        translated_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
         return translated_text
         
     except Exception as e:
         print(f"Detailed error: {str(e)}")
         return f"Error during translation: {str(e)}"
 
+# Create and launch the interface
+demo = gr.Interface(
+    fn=gradio_predict,
+    inputs=gr.Textbox(label="English text"),
+    outputs=gr.Textbox(label="French translation"),
+    title="English to French Translator",
+    description="Enter English text to translate to French"
+)
+
+if __name__ == "__main__":
+    demo.launch()
+
 # Gradio interface for the web app
 gr.Interface(
     fn=gradio_predict,