Init app

app.py

@@ -3,19 +3,23 @@ from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel, PeftConfig
 import torch
-# Initialize the FastAPI app
+import bitsandbytes as bnb
+
+# Initialize FastAPI app
 app = FastAPI()
 
-# Load model and tokenizer once at startup
+# Load model and tokenizer in optimized mode at startup
 base_model_name = "akjindal53244/Llama-3.1-Storm-8B"
 peft_model_id = "LlamaFactoryAI/cv-job-description-matching"
 
-base_model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.float16)
-model = PeftModel.from_pretrained(base_model, peft_model_id, torch_dtype=torch.float16)
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_name,
+    torch_dtype=torch.float16,     # Use FP16 precision
+    load_in_8bit=True,             # Enable Int8 quantization
+    device_map="auto"              # Auto-manage device allocation
+)
 tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-config = PeftConfig.from_pretrained(peft_model_id)
-
-
+model = PeftModel.from_pretrained(base_model, peft_model_id, torch_dtype=torch.float16)
 
 # Define request model
 class AnalysisRequest(BaseModel):
@@ -27,31 +31,19 @@ async def analyze(request: AnalysisRequest):
     try:
         # Prepare input text with formatted message
         system_prompt = """
-        You are an advanced AI model designed to analyze the compatibility between a CV and a job description. You will receive a CV and a job description. Your task is to output a structured JSON format that includes the following:
-
-        1. matching_analysis: Analyze the CV against the job description to identify key strengths and gaps.
-        2. description: Summarize the relevance of the CV to the job description in a few concise sentences.
-        3. score: Provide a numerical compatibility score (0-100) based on qualifications, skills, and experience.
-        4. recommendation: Suggest actions for the candidate to improve their match or readiness for the role.
-
-        Your output must be in JSON format as follows:
-        {
-          "matching_analysis": "Your detailed analysis here.",
-          "description": "A brief summary here.",
-          "score": 85,
-          "recommendation": "Your suggestions here."
-        }
+        You are an advanced AI model designed to analyze the compatibility between a CV and a job description...
         """
         user_input = f"<CV> {request.cv} </CV>\n<job_description> {request.job_description} </job_description>"
         input_text = system_prompt + user_input
 
-        # Tokenize and generate response
-        inputs = tokenizer(input_text, return_tensors="pt")
-        outputs = model.generate(**inputs, max_new_tokens=64)
+        # Tokenize and generate response with memory optimizations
+        inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
+        
+        with torch.no_grad():  # Disable gradients for inference
+            outputs = model.generate(**inputs, max_new_tokens=64)  # Limit generated tokens
         generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
         return {"analysis": generated_text}
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
-