updated

app.py

@@ -21,8 +21,9 @@ base_model = AutoModelForCausalLM.from_pretrained(
 # Load adapter and ensure it is on the correct device
 model = PeftModel.from_pretrained(base_model, adapter_path).to(device)
 
-# Load tokenizer
+# Load tokenizer and ensure padding token is set
 tokenizer = AutoTokenizer.from_pretrained(base_model_path)
+tokenizer.pad_token = tokenizer.eos_token  # Avoids padding issues
 
 # Define request model for validation
 class GenerateRequest(BaseModel):
@@ -40,6 +41,7 @@ def generate_text_from_model(prompt: str):
         input_ids = input_data.input_ids.to(device)
         attention_mask = input_data.attention_mask.to(device)
 
+        # Generate output
         output_ids = model.generate(
             input_ids, 
             max_length=512, 
@@ -52,6 +54,9 @@ def generate_text_from_model(prompt: str):
         # Extract only the assistant's response
         response_text = generated_text.split("<|assistant|>\n")[-1].strip()
         return response_text
+    except torch.cuda.OutOfMemoryError:
+        torch.cuda.empty_cache()
+        raise HTTPException(status_code=500, detail="CUDA Out of Memory. Try using a smaller model or lowering max_length.")
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -65,4 +70,4 @@ async def root():
 @app.post("/generate/")
 async def generate_text(request: GenerateRequest):
     response = generate_text_from_model(request.prompt)
-    return {"response": response}
+    return {"response": response}