Update functions.py

functions.py

@@ -6,7 +6,7 @@ from bs4 import BeautifulSoup
 from peft import PeftConfig, PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
 
 generation_config = GenerationConfig(temperature=.8,
@@ -46,7 +46,7 @@ def load_model(peft_model_id):
     config = PeftConfig.from_pretrained(peft_model_id)
     print(['load_model', 'loading model'])
     model = AutoModelForCausalLM.from_pretrained(
-        config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map='auto')
+        config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map='cuda')
     print(['load_model', 'loading tokenizer'])
     tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
     model = PeftModel.from_pretrained(model, peft_model_id)