Update app.py

app.py

@@ -10,17 +10,17 @@ logger = logging.getLogger(__name__)
 
 app = FastAPI(title="SQL Coder API")
 
-# Ensure cache directory exists
-cache_dir = os.getenv('TRANSFORMERS_CACHE', '/home/user/.cache/huggingface')
-os.makedirs(cache_dir, exist_ok=True)
+# Set environment variable for cache directory
+os.environ['TRANSFORMERS_CACHE'] = '/home/user/.cache/huggingface'
 
 # Initialize pipeline
 try:
-    pipe = pipeline("text-generation", 
-                   model="defog/llama-3-sqlcoder-8b",
-                   device_map="auto",
-                   torch_dtype="auto",
-                   cache_dir=cache_dir)
+    pipe = pipeline(
+        "text-generation",
+        model="defog/llama-3-sqlcoder-8b",
+        device_map="auto",
+        model_kwargs={"torch_dtype": "auto"}
+    )
     logger.info("Pipeline initialized successfully")
 except Exception as e:
     logger.error(f"Error initializing pipeline: {str(e)}")
@@ -32,7 +32,7 @@ class ChatMessage(BaseModel):
 
 class QueryRequest(BaseModel):
     messages: list[ChatMessage]
-    max_length: int = 1024
+    max_new_tokens: int = 1024
     temperature: float = 0.7
 
 class QueryResponse(BaseModel):
@@ -47,10 +47,11 @@ async def generate(request: QueryRequest):
         # Generate response using pipeline
         response = pipe(
             formatted_prompt,
-            max_length=request.max_length,
+            max_new_tokens=request.max_new_tokens,
             temperature=request.temperature,
             do_sample=True,
-            num_return_sequences=1
+            num_return_sequences=1,
+            pad_token_id=pipe.tokenizer.eos_token_id
         )
         
         # Extract generated text