updated

app.py

@@ -31,13 +31,31 @@ class GenerateRequest(BaseModel):
 # **Use model.generate() instead of pipeline()**
 def generate_text_from_model(prompt: str):
     try:
-        input_ids = tokenizer(f"<s>[INST] {prompt} [/INST]", return_tensors="pt").input_ids.to(device)
-        output_ids = model.generate(input_ids, max_length=512)
+        input_data = tokenizer(
+            f"<s>[INST] {prompt} [/INST]", 
+            return_tensors="pt", 
+            padding=True, 
+            truncation=True
+        )
+        input_ids = input_data.input_ids.to(device)
+        attention_mask = input_data.attention_mask.to(device)
+
+        output_ids = model.generate(
+            input_ids, 
+            max_length=512, 
+            pad_token_id=tokenizer.eos_token_id,  # Explicitly setting pad_token_id
+            attention_mask=attention_mask
+        )
+
         generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-        return generated_text
+
+        # Extract only the assistant's response
+        response_text = generated_text.split("<|assistant|>\n")[-1].strip()
+        return response_text
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+
 # Root endpoint for testing
 @app.get("/")
 async def root():