Upload app.py

optimize app.py

app.py

@@ -23,6 +23,10 @@ model = AutoModelForCausalLM.from_pretrained(
     token=HF_API_TOKEN
 )
 
+# Ensure tokenizer has a padding token
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token  # Use the EOS token as the padding token
+
 @app.post("/generate")
 async def generate_text(request: Request):
     data = await request.json()
@@ -31,15 +35,15 @@ async def generate_text(request: Request):
 
     # Tokenize with padding and attention mask
     inputs = tokenizer(
-        prompt, 
-        return_tensors="pt", 
-        padding=True, 
+        prompt,
+        return_tensors="pt",
+        padding=True,
         truncation=True
     ).to(model.device)
 
     outputs = model.generate(
-        inputs["input_ids"], 
-        attention_mask=inputs["attention_mask"], 
+        inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
         max_length=max_length
     )