Update app.py

app.py

@@ -1,28 +1,69 @@
 import os
-
-
-from fastapi import FastAPI
-from pydantic import BaseModel
+import logging
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from typing import Optional
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # 1. Load model & tokenizer once at startup
 MODEL_ID = "EQuIP-Queries/EQuIP_3B"
-# Specify cache_dir just in case
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+except Exception as e:
+    logger.error(f"Failed to load model: {e}")
+    raise
 
 # 2. Initialize FastAPI
-app = FastAPI()
+app = FastAPI(title="EQuIP Query Generator",
+             description="Generate Elasticsearch queries using EQuIP model",
+             version="1.0.0")
 
-# 3. Define request schema
+# 3. Define request/response schemas
 class GenerateRequest(BaseModel):
-    prompt: str
-    max_new_tokens: int = 50
+    prompt: str = Field(..., description="Input prompt for query generation")
+    max_new_tokens: int = Field(default=50, ge=1, le=512, description="Maximum number of tokens to generate")
+
+class GenerateResponse(BaseModel):
+    generated_text: str
+    input_prompt: str
+    token_count: Optional[int]
+
+# 4. Health check endpoint
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy", "model": MODEL_ID}
 
-# 4. Inference endpoint
-@app.post("/generate")
+# 5. Inference endpoint
+@app.post("/generate", response_model=GenerateResponse)
 async def generate(req: GenerateRequest):
-    inputs = tokenizer(req.prompt, return_tensors="pt")
-    ids = model.generate(**inputs, max_new_tokens=req.max_new_tokens)
-    text = tokenizer.decode(ids[0], skip_special_tokens=True)
-    return {"generated_text": text}
+    try:
+        logger.info(f"Processing request with prompt: {req.prompt[:50]}...")
+        inputs = tokenizer(req.prompt, return_tensors="pt")
+        
+        ids = model.generate(
+            **inputs,
+            max_new_tokens=req.max_new_tokens,
+            pad_token_id=tokenizer.eos_token_id,
+            num_return_sequences=1
+        )
+        
+        generated_text = tokenizer.decode(ids[0], skip_special_tokens=True)
+        token_count = len(ids[0])
+        
+        return GenerateResponse(
+            generated_text=generated_text,
+            input_prompt=req.prompt,
+            token_count=token_count
+        )
+        
+    except Exception as e:
+        logger.error(f"Generation failed: {str(e)}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Generation failed: {str(e)}"
+        )