Fix main.py: add os import, health endpoint, memory management, use instruction

main.py

@@ -15,51 +15,51 @@ app = FastAPI(title="LaMini-LM API",
               description="API for text generation using LaMini-GPT-774M", version="1.0.0")
 
 # Define request model
-
-
 class TextGenerationRequest(BaseModel):
-    prompt: str
+    instruction: str  # Changed from prompt for consistency
     max_length: int = 100
     temperature: float = 1.0
     top_p: float = 0.9
 
-
 # Load model (cached after first load)
 try:
     logger.info("Loading LaMini-GPT-774M model...")
-    # device=-1 for CPU
-    generator = pipeline(
-        'text-generation', model='MBZUAI/LaMini-GPT-774M', device=-1)
+    generator = pipeline('text-generation', model='MBZUAI/LaMini-GPT-774M', device=-1)
     logger.info("Model loaded successfully.")
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
 except Exception as e:
     logger.error(f"Failed to load model: {str(e)}")
-    raise Exception(f"Model loading failed: {str(e)}")
+    generator = None  # Allow server to run for health check
 
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}
+
+@app.get("/")
+async def root():
+    return {"message": "Welcome to the LaMini-LM API. Use POST /generate to generate text."}
 
 @app.post("/generate")
 async def generate_text(request: TextGenerationRequest):
-    """
-    Generate text based on the input prompt using LaMini-GPT-774M.
-    """
+    if generator is None:
+        raise HTTPException(status_code=503, detail="Model not loaded. Check server logs.")
     try:
         # Validate inputs
-        if not request.prompt.strip():
-            raise HTTPException(
-                status_code=400, detail="Prompt cannot be empty")
+        if not request.instruction.strip():
+            raise HTTPException(status_code=400, detail="Instruction cannot be empty")
         if request.max_length < 10 or request.max_length > 500:
-            raise HTTPException(
-                status_code=400, detail="max_length must be between 10 and 500")
+            raise HTTPException(status_code=400, detail="max_length must be between 10 and 500")
         if request.temperature <= 0 or request.temperature > 2:
-            raise HTTPException(
-                status_code=400, detail="temperature must be between 0 and 2")
+            raise HTTPException(status_code=400, detail="temperature must be between 0 and 2")
         if request.top_p <= 0 or request.top_p > 1:
-            raise HTTPException(
-                status_code=400, detail="top_p must be between 0 and 1")
+            raise HTTPException(status_code=400, detail="top_p must be between 0 and 1")
 
         # Generate text
-        logger.info(f"Generating text for prompt: {request.prompt[:50]}...")
-        wrapper = "Instruction: You are a helpful assistant. Please respond to the following prompt.\n\nPrompt: {}\n\nResponse:".format(
-            request.prompt)
+        logger.info(f"Generating text for instruction: {request.instruction[:50]}...")
+        wrapper = "Instruction: You are a helpful assistant. Please respond to the following instruction.\n\nInstruction: {}\n\nResponse:".format(
+            request.instruction)
         outputs = generator(
             wrapper,
             max_length=request.max_length,
@@ -68,19 +68,13 @@ async def generate_text(request: TextGenerationRequest):
             num_return_sequences=1,
             do_sample=True
         )
-        generated_text = outputs[0]['generated_text'].replace(
-            wrapper, "").strip()
-
+        generated_text = outputs[0]['generated_text'].replace(wrapper, "").strip()
         return {"generated_text": generated_text}
     except Exception as e:
         logger.error(f"Error during text generation: {str(e)}")
-        raise HTTPException(
-            status_code=500, detail=f"Text generation failed: {str(e)}")
-
+        raise HTTPException(status_code=500, detail=f"Text generation failed: {str(e)}")
 
-@app.get("/")
-async def root():
-    """
-    Root endpoint with basic info.
-    """
-    return {"message": "Welcome to the LaMini-LM API. Use POST /generate to generate text."}
+if __name__ == "__main__":
+    import uvicorn
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run(app, host="0.0.0.0", port=port)