Update model.py

model.py

@@ -1,8 +1,21 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import os
 import torch
-if torch.cuda.is_available():
-    model.to('cuda')
+
+# Check if CUDA is available for faster inference
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# 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
+)
+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):
     """
     Generate ATS feedback by utilizing a pre-configured pipeline.
@@ -32,28 +45,19 @@ def modelFeedback(ats_score, resume_data, job_description):
         #### Job Description: {job_description}
     """
 
-    # Load the tokenizer and model
-    huggingface_token = os.environ.get("KEY2")
-    tokenizer = AutoTokenizer.from_pretrained(
-        "meta-llama/Llama-3.2-1B",
-        use_auth_token=huggingface_token
-    )
-    model = AutoModelForCausalLM.from_pretrained(
-        "meta-llama/Llama-3.2-1B",
-        use_auth_token=huggingface_token
-    )
-
     try:
         # Tokenize the input
-        input_ids = tokenizer.encode(input_prompt, return_tensors="pt")
+        input_ids = tokenizer.encode(input_prompt, return_tensors="pt").to(device)
         
-        # 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
-        )
+        # 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)