Update model.py

model.py

@@ -1,20 +1,16 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import os
 import torch
+from transformers import pipeline
 
-# Check if CUDA is available for faster inference
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+# Define model id
+model_id = "meta-llama/Llama-3.2-1B"
 
-# Load the tokenizer and model once, outside of the function
-huggingface_token = os.environ.get("KEY2")
-tokenizer = AutoTokenizer.from_pretrained(
-    "meta-llama/Llama-3.2-1B",
-    use_auth_token=huggingface_token
+# Create pipeline for text generation with bfloat16 precision and device auto-placement
+pipe = pipeline(
+    "text-generation", 
+    model=model_id, 
+    torch_dtype=torch.bfloat16, 
+    device_map="auto"
 )
-model = AutoModelForCausalLM.from_pretrained(
-    "meta-llama/Llama-3.2-1B",
-    use_auth_token=huggingface_token
-).to(device)
 
 def modelFeedback(ats_score, resume_data, job_description):
     """
@@ -46,21 +42,11 @@ def modelFeedback(ats_score, resume_data, job_description):
     """
 
     try:
-        # Tokenize the input
-        input_ids = tokenizer.encode(input_prompt, return_tensors="pt").to(device)
+        # Generate the feedback using the pre-configured pipeline
+        response = pipe(input_prompt, max_length=1500, num_return_sequences=1)
         
-        # Disable gradient calculation for faster inference
-        with torch.no_grad():
-            # Generate the output
-            output = model.generate(
-                input_ids,
-                max_length=1500,
-                temperature=0.01,
-                pad_token_id=tokenizer.eos_token_id  # Ensure padding works properly
-            )
-        
-        # Decode the output
-        response_text = tokenizer.decode(output[0], skip_special_tokens=True)
+        # Extract the generated text
+        response_text = response[0]['generated_text']
         return response_text
     except Exception as e:
         print(f"Error during generation: {e}")