Update app.py

app.py

@@ -1,16 +1,31 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 from huggingface_hub import snapshot_download
 from safetensors.torch import load_file
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 class ModelInput(BaseModel):
-    prompt: str
-    max_new_tokens: int = 2048
+    prompt: str = Field(..., description="The input prompt for text generation")
+    max_new_tokens: int = Field(default=2048, gt=0, le=4096, description="Maximum number of tokens to generate")
 
 app = FastAPI()
 
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
 # Define model paths
 BASE_MODEL_PATH = "HuggingFaceTB/SmolLM2-135M-Instruct"
 ADAPTER_PATH = "khurrameycon/SmolLM-135M-Instruct-qa_pairs_converted.json-25epochs"
@@ -18,7 +33,7 @@ ADAPTER_PATH = "khurrameycon/SmolLM-135M-Instruct-qa_pairs_converted.json-25epoc
 def load_model_and_tokenizer():
     """Load the model, tokenizer, and adapter weights."""
     try:
-        print("Loading base model...")
+        logger.info("Loading base model...")
         model = AutoModelForCausalLM.from_pretrained(
             BASE_MODEL_PATH,
             torch_dtype=torch.float16,
@@ -26,31 +41,38 @@ def load_model_and_tokenizer():
             device_map="auto"
         )
 
-        print("Loading tokenizer...")
+        logger.info("Loading tokenizer...")
         tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_PATH)
 
-        print("Downloading adapter weights...")
+        logger.info("Downloading adapter weights...")
         adapter_path_local = snapshot_download(repo_id=ADAPTER_PATH)
 
-        print("Loading adapter weights...")
+        logger.info("Loading adapter weights...")
         adapter_file = f"{adapter_path_local}/adapter_model.safetensors"
         state_dict = load_file(adapter_file)
 
-        print("Applying adapter weights...")
+        logger.info("Applying adapter weights...")
         model.load_state_dict(state_dict, strict=False)
-        print("Model and adapter loaded successfully!")
+        logger.info("Model and adapter loaded successfully!")
 
         return model, tokenizer
     except Exception as e:
-        print(f"Error during model loading: {e}")
+        logger.error(f"Error during model loading: {e}", exc_info=True)
         raise
 
 # Load model and tokenizer at startup
-model, tokenizer = load_model_and_tokenizer()
+try:
+    model, tokenizer = load_model_and_tokenizer()
+except Exception as e:
+    logger.error(f"Failed to load model at startup: {e}", exc_info=True)
+    model = None
+    tokenizer = None
 
 def generate_response(model, tokenizer, instruction, max_new_tokens=2048):
     """Generate a response from the model based on an instruction."""
     try:
+        logger.info(f"Generating response for instruction: {instruction[:100]}...")
+        
         # Encode input with truncation
         inputs = tokenizer.encode(
             instruction,
@@ -59,6 +81,8 @@ def generate_response(model, tokenizer, instruction, max_new_tokens=2048):
             max_length=tokenizer.model_max_length
         ).to(model.device)
 
+        logger.info(f"Input shape: {inputs.shape}")
+        
         # Create attention mask
         attention_mask = torch.ones(inputs.shape, device=model.device)
 
@@ -70,35 +94,59 @@ def generate_response(model, tokenizer, instruction, max_new_tokens=2048):
             temperature=0.7,
             top_p=0.9,
             do_sample=True,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
         )
 
+        logger.info(f"Output shape: {outputs.shape}")
+
         # Decode and strip input prompt from response
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         generated_text = response[len(instruction):].strip()
 
+        logger.info(f"Generated text length: {len(generated_text)}")
         return generated_text
     except Exception as e:
-        print(f"Error generating response: {e}")
+        logger.error(f"Error generating response: {e}", exc_info=True)
         raise ValueError(f"Error generating response: {e}")
 
 @app.post("/generate")
-async def generate_text(input: ModelInput):
+async def generate_text(input: ModelInput, request: Request):
     """Generate text based on the input prompt."""
     try:
-        print(f"Received prompt: {input.prompt}")
+        if model is None or tokenizer is None:
+            raise HTTPException(status_code=503, detail="Model not loaded")
+
+        logger.info(f"Received request from {request.client.host}")
+        logger.info(f"Prompt: {input.prompt[:100]}...")
+        
         response = generate_response(
             model=model,
             tokenizer=tokenizer,
             instruction=input.prompt,
             max_new_tokens=input.max_new_tokens
         )
-        print(f"Generated response: {response}")
+        
+        if not response:
+            logger.warning("Generated empty response")
+            return {"generated_text": "", "warning": "Empty response generated"}
+            
+        logger.info(f"Generated response length: {len(response)}")
         return {"generated_text": response}
     except Exception as e:
-        print(f"Error: {str(e)}")
+        logger.error(f"Error in generate_text endpoint: {e}", exc_info=True)
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/")
 async def root():
     """Root endpoint that returns a welcome message."""
-    return {"message": "Welcome to the Model API!"}
+    return {"message": "Welcome to the Model API!", "status": "running"}
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint."""
+    return {
+        "status": "healthy",
+        "model_loaded": model is not None and tokenizer is not None,
+        "model_device": str(next(model.parameters()).device) if model else None
+    }