add cuda

app.py

@@ -25,7 +25,7 @@ logger = logging.getLogger(__name__)
 login(token=os.environ['HF_TOKEN'])
 
 # Global variables
-CUR_MODEL = "./repllama-instruct-hard-positives-v2-joint-full-weights"
+CUR_MODEL = "./repllama-instruct-hard-positives-v2-joint"
 BASE_MODEL = "meta-llama/Llama-2-7b-hf"
 tokenizer = None
 model = None
@@ -62,15 +62,15 @@ def create_batch_dict(tokenizer, input_texts, max_length=512):
     )
 
 def load_model():
-    global tokenizer, model, CUR_MODEL
+    global tokenizer, model, CUR_MODEL, BASE_MODEL
     tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
     tokenizer.pad_token_id = tokenizer.eos_token_id
     tokenizer.pad_token = tokenizer.eos_token
     tokenizer.padding_side = "right"
     
-    model = AutoModel.from_pretrained(CUR_MODEL, max_memory={"cpu": "12GiB"}, torch_dtype=torch.bfloat16, offload_state_dict=True)
-    # base_model_instance = AutoModel.from_pretrained(BASE_MODEL)
-    # model = PeftModel.from_pretrained(base_model_instance, CUR_MODEL)
+    # model = AutoModel.from_pretrained(CUR_MODEL, max_memory={"cpu": "12GiB"}, torch_dtype=torch.bfloat16, offload_state_dict=True)
+    base_model_instance = AutoModel.from_pretrained(BASE_MODEL)
+    model = PeftModel.from_pretrained(base_model_instance, CUR_MODEL)
     model.eval()