use accelerators in app

src/models/llm_manager.py

@@ -14,12 +14,26 @@ class LLMManager:
     def __init__(self):
         self.model = None
         self.tokenizer = None
-        
+
+        if torch.cuda.is_available():
+            device = "cuda"
+            dtype = torch.float16
+        elif torch.backends.mps.is_available():
+            device = "mps"
+            dtype = torch.float16
+        else:
+            device = "cpu"
+            dtype = torch.float32
+
+        self.device = device
+        self.dtype = dtype
+
     def load_models(self, model_name="meta-llama/Llama-3.2-1B-Instruct"):
         """Load the LLM model and tokenizer"""
         self.tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)
-        self.model = AutoModelForCausalLM.from_pretrained(model_name, token=HF_TOKEN)
-        
+        self.model = AutoModelForCausalLM.from_pretrained(model_name, dtype=self.dtype, low_cpu_mem_usage=True)
+        self.model = self.model.to(self.device)
+
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
     
@@ -60,7 +74,8 @@ class LLMManager:
         full_prompt = f"{prompt}\n\nAssistant: {partial_response}{user_continuation}"
         
         inputs = self.tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True)
-        
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
         with torch.no_grad():
             outputs = self.model.generate(
                 inputs.input_ids,