Update app.py

app.py

@@ -38,19 +38,28 @@ class LocalBartModel:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model.to(self.device)
 
-    def generate(self, input_ids, **generate_kwargs):
-        # Defensive: convert list input_ids to tensor if needed
-        if isinstance(input_ids, list):
-            input_ids = torch.tensor(input_ids)
+    def generate(self, inputs, **generate_kwargs):
+        input_ids = inputs.get("input_ids")
+        attention_mask = inputs.get("attention_mask")
+
+        if input_ids is None:
+            raise ValueError("input_ids missing from tokenizer output")
+
         input_ids = input_ids.to(self.device)
-        return self.model.generate(input_ids, **generate_kwargs)
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(self.device)
+
+        return self.model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            **generate_kwargs
+        )
 
     def __call__(self, prompt: str) -> str:
         inputs = self.tokenizer(prompt, return_tensors="pt")
-        input_ids = inputs.input_ids  # tensor here
 
         output_ids = self.generate(
-            input_ids,
+            inputs,
             max_length=100,
             num_beams=5,
             early_stopping=True