Update main.py

main.py

@@ -19,11 +19,17 @@ model = AutoModelForCausalLM.from_pretrained(
 
 # Example usage: Generate text
 prompt = "The quick brown fox"
-input_ids = tokenizer.encode(prompt, return_tensors="pt")
-output = model.generate(input_ids, max_length=50, num_return_sequences=1)
-generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
 
-print(generated_text)
+inputs = tokenizer(prompt, return_tensors="pt", padding=True, return_attention_mask=True, ).to(model.device)
+outputs = model.generate(
+    **inputs,
+    max_new_tokens=100,
+    pad_token_id=tokenizer.eos_token_id  # Set this to suppress warning
+)
+
+resp = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+print(resp)
 
 
 app = FastAPI()
@@ -53,7 +59,7 @@ async def generate_text(item: Item):
 
     # logging.info("Response generated")
 
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True, return_attention_mask=True, ).to(model.device)
+    inputs = tokenizer(item.prompt, return_tensors="pt", padding=True, return_attention_mask=True, ).to(model.device)
 
 
     # input_ids = tokenizer.encode(item.prompt, return_tensors="pt")