Upload app.py

app.py

@@ -29,8 +29,19 @@ async def generate_text(request: Request):
     prompt = data.get("prompt", "")
     max_length = data.get("max_length", 50)
 
-    # Tokenize and generate
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(inputs["input_ids"], max_length=max_length)
+    # Tokenize with padding and attention mask
+    inputs = tokenizer(
+        prompt, 
+        return_tensors="pt", 
+        padding=True, 
+        truncation=True
+    ).to(model.device)
+
+    outputs = model.generate(
+        inputs["input_ids"], 
+        attention_mask=inputs["attention_mask"], 
+        max_length=max_length
+    )
+
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return {"generated_text": response}