Update app.py

app.py

@@ -10,9 +10,9 @@ import logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-app = FastAPI(title="llm-apiku", version="1.0.0")
+app = FastAPI(title="Gema 4B Model API", version="1.0.0")
 
-# Request model - fleksibel untuk menerima semua parameter
+# Request model
 class TextRequest(BaseModel):
     inputs: str
     system_prompt: Optional[str] = None
@@ -33,34 +33,43 @@ model = None
 @app.on_event("startup")
 async def load_model():
     global model
+    # Define the local model path
+    model_path = "./model"
+    model_file = "gema-4b-indra10k-model1-q4_k_m.gguf"
+    
     try:
-        logger.info("Loading model...")
+        if not os.path.exists(model_path) or not os.path.exists(os.path.join(model_path, model_file)):
+             raise RuntimeError("Model files not found. Ensure the model was downloaded in the Docker build.")
+
+        logger.info(f"Loading model from local path: {model_path}")
+        # Load the model from the local directory downloaded during the Docker build
         model = AutoModelForCausalLM.from_pretrained(
-            "Dnfs/gema-4b-indra10k-model1-Q4_K_M-GGUF",
-            model_file="gema-4b-indra10k-model1-q4_k_m.gguf",
+            model_path, # Load from the local folder
+            model_file=model_file, # Specify the GGUF file name
             model_type="llama",
-            gpu_layers=0,  # Set to appropriate number if using GPU
+            gpu_layers=0,
             context_length=2048,
-            threads=os.cpu_count()
+            threads=os.cpu_count() or 1
         )
         logger.info("Model loaded successfully!")
     except Exception as e:
         logger.error(f"Failed to load model: {e}")
+        # Raising the exception will prevent the app from starting if the model fails to load
         raise e
 
 @app.post("/generate", response_model=TextResponse)
 async def generate_text(request: TextRequest):
     if model is None:
-        raise HTTPException(status_code=500, detail="Model not loaded")
+        raise HTTPException(status_code=503, detail="Model is not ready or failed to load. Please try again later.")
     
     try:
-        # Buat prompt - gunakan system_prompt jika ada, atau langsung input user
+        # Create prompt
         if request.system_prompt:
             full_prompt = f"{request.system_prompt}\n\nUser: {request.inputs}\nAssistant:"
         else:
             full_prompt = request.inputs
         
-        # Generate text dengan parameter dari request
+        # Generate text with parameters from the request
         generated_text = model(
             full_prompt,
             max_new_tokens=request.max_tokens,
@@ -71,7 +80,7 @@ async def generate_text(request: TextRequest):
             stop=request.stop or []
         )
         
-        # Bersihkan response dari system prompt jika ada
+        # Clean up the response
         if "Assistant:" in generated_text:
             generated_text = generated_text.split("Assistant:")[-1].strip()
         
@@ -83,6 +92,8 @@ async def generate_text(request: TextRequest):
 
 @app.get("/health")
 async def health_check():
+    # The health check now also implicitly checks if the model has been loaded
+    # because a failure in load_model will stop the app from running.
     return {"status": "healthy", "model_loaded": model is not None}
 
 @app.get("/")