Update app.py

app.py

@@ -31,7 +31,6 @@ SYSTEM_PROMPT = (
     "Never say 'the answer is...'. Only return the answer.\n"
 )
 
-# Local wrapper for facebook/bart-base that exposes generate()
 class LocalBartModel:
     def __init__(self, model_name="facebook/bart-base"):
         self.tokenizer = BartTokenizer.from_pretrained(model_name)
@@ -40,12 +39,18 @@ class LocalBartModel:
         self.model.to(self.device)
 
     def generate(self, input_ids, **generate_kwargs):
-        return self.model.generate(input_ids.to(self.device), **generate_kwargs)
+        # Defensive: convert list input_ids to tensor if needed
+        if isinstance(input_ids, list):
+            input_ids = torch.tensor(input_ids)
+        input_ids = input_ids.to(self.device)
+        return self.model.generate(input_ids, **generate_kwargs)
 
     def __call__(self, prompt: str) -> str:
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        inputs = self.tokenizer(prompt, return_tensors="pt")
+        input_ids = inputs.input_ids  # tensor here
+
         output_ids = self.generate(
-            inputs.input_ids,
+            input_ids,
             max_length=100,
             num_beams=5,
             early_stopping=True