check model location

agents/mistral_agent.py

@@ -42,7 +42,7 @@ class MistralAgent(object):
         self.padding_side = "right"
         self.model = AutoModelForCausalLM.from_pretrained("./agents/mistral_model/model",
                                                           quantization_config=bnb_config,
-                                                          device_map=device_map)
+                                                          device_map="auto")
 
         self.model.config.use_cache = True # For faster generation
 
@@ -76,6 +76,7 @@ class MistralAgent(object):
         # self.sep_id = self.tokenizer.sep_token_id
 
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        print(self.device)
 
         # self.model = self.model.to(self.device)
         self.model.eval()

requirements.txt

@@ -11,3 +11,4 @@ scipy
 sentencepiece
 torch
 tqdm
+protobuf