Update main.py

main.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException, Request
+from fastapi import FastAPI, File, UploadFile, HTTPException, Request, BackgroundTasks
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 import pandas as pd
@@ -9,23 +9,21 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import uvicorn
 import nest_asyncio
 import os
+import uuid
 
-# Allow uvicorn to run inside Colab (safe for Spaces / Render too)
+# Allow uvicorn to run inside Colab / Spaces
 nest_asyncio.apply()
+
 # ========== Environment Setup ==========
 device = "cpu"
-
-# Set offload folder path
 offload_dir = "/data/model_offload"
 os.makedirs(offload_dir, exist_ok=True)
 
 adapter_path = "./phi3-ds-expert-lora"
 base_model_id = "microsoft/Phi-3-mini-4k-instruct"
 
-# Load LoRA config
 peft_config = PeftConfig.from_pretrained(adapter_path)
 
-# Load base model
 base_model = AutoModelForCausalLM.from_pretrained(
     base_model_id,
     torch_dtype="auto",
@@ -36,7 +34,6 @@ base_model = AutoModelForCausalLM.from_pretrained(
 tokenizer = AutoTokenizer.from_pretrained(base_model_id)
 tokenizer.pad_token = tokenizer.eos_token
 
-# Load LoRA adapter with explicit offload_dir
 model = PeftModel.from_pretrained(
     base_model,
     adapter_path,
@@ -46,7 +43,6 @@ model = PeftModel.from_pretrained(
 model.to(device)
 model.eval()
 
-
 # ========== FastAPI App ==========
 app = FastAPI()
 
@@ -62,6 +58,9 @@ app.add_middleware(
 # Global dataset storage
 uploaded_df = None
 
+# In-memory job results store
+job_results = {}
+
 # ========== Pydantic Request Schema ==========
 class PromptRequest(BaseModel):
     prompt: str
@@ -85,18 +84,22 @@ def generate_prompt(prompt_text, df=None):
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=128,  # fast generation
-            do_sample=False,     # deterministic
+            max_new_tokens=128,
+            do_sample=False,
             eos_token_id=tokenizer.eos_token_id,
             pad_token_id=tokenizer.eos_token_id
         )
 
-    # Decode only the new generated tokens (ignoring prompt)
     generated_tokens = outputs[0][input_length:]
     response = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
 
     return response
 
+# ========== Background Inference Task ==========
+def run_inference(prompt_text, job_id, df=None):
+    result = generate_prompt(prompt_text, df)
+    job_results[job_id] = result
+
 # ========== API Endpoints ==========
 
 @app.get("/")
@@ -104,24 +107,37 @@ def read_root():
     return {"message": "🚀 API is live and running!"}
 
 @app.post("/generate")
-async def generate(request: PromptRequest):
+async def generate(request: PromptRequest, background_tasks: BackgroundTasks):
     if not request.prompt.strip():
         raise HTTPException(status_code=400, detail="Prompt cannot be empty.")
-    result = generate_prompt(request.prompt, df=uploaded_df)
-    return {"response": result}
+
+    job_id = str(uuid.uuid4())
+    job_results[job_id] = "processing"
+    background_tasks.add_task(run_inference, request.prompt, job_id, uploaded_df)
+    return {"job_id": job_id}
+
+@app.get("/status/{job_id}")
+async def job_status(job_id: str):
+    if job_id not in job_results:
+        raise HTTPException(status_code=404, detail="Job ID not found")
+
+    result = job_results[job_id]
+    if result == "processing":
+        return {"status": "processing"}
+    else:
+        return {"status": "complete", "result": result}
 
 MAX_FILE_SIZE = 15 * 1024 * 1024  # 15 MB
 
 @app.post("/upload")
 async def upload_file(request: Request, file: UploadFile = File(...)):
-    # Check content-length header if available
+    global uploaded_df
     content_length = request.headers.get("content-length")
     if content_length and int(content_length) > MAX_FILE_SIZE:
         raise HTTPException(status_code=413, detail="File too large. Max size is 15 MB.")
 
     content = await file.read()
 
-    # Double-check actual content size after reading
     if len(content) > MAX_FILE_SIZE:
         raise HTTPException(status_code=413, detail="File too large. Max size is 15 MB.")
 
@@ -132,9 +148,6 @@ async def upload_file(request: Request, file: UploadFile = File(...)):
             uploaded_df = pd.read_excel(io.BytesIO(content))
         else:
             raise HTTPException(status_code=400, detail="Only CSV and Excel files are supported.")
-
-        
-
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error while uploading file: {str(e)}")