5to9 commited on
Commit
b235bfd
1 Parent(s): d0aacc5

0.14 explictly set GPU and dtype

Browse files
Files changed (1) hide show
  1. app.py +4 -4
app.py CHANGED
@@ -24,11 +24,11 @@ models_available = [
24
  "mistralai/Mistral-7B-Instruct-v0.3",
25
  ]
26
 
27
- device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
28
 
29
  tokenizer_a, model_a = None, None
30
  tokenizer_b, model_b = None, None
31
- torch_dtype = torch.bfloat16
32
 
33
  def apply_chat_template(messages, add_generation_prompt=False):
34
  """
@@ -62,7 +62,7 @@ def load_model_a(model_id):
62
  tokenizer_a = AutoTokenizer.from_pretrained(model_id)
63
  model_a = AutoModelForCausalLM.from_pretrained(
64
  model_id,
65
- torch_dtype=torch_dtype,
66
  device_map="auto",
67
  trust_remote_code=True,
68
  ).eval()
@@ -79,7 +79,7 @@ def load_model_b(model_id):
79
  tokenizer_b = AutoTokenizer.from_pretrained(model_id)
80
  model_b = AutoModelForCausalLM.from_pretrained(
81
  model_id,
82
- torch_dtype=torch_dtype,
83
  device_map="auto",
84
  trust_remote_code=True,
85
  ).eval()
 
24
  "mistralai/Mistral-7B-Instruct-v0.3",
25
  ]
26
 
27
+ #device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
28
+ device = "cuda"
29
 
30
  tokenizer_a, model_a = None, None
31
  tokenizer_b, model_b = None, None
 
32
 
33
  def apply_chat_template(messages, add_generation_prompt=False):
34
  """
 
62
  tokenizer_a = AutoTokenizer.from_pretrained(model_id)
63
  model_a = AutoModelForCausalLM.from_pretrained(
64
  model_id,
65
+ torch_dtype=torch.bfloat16,
66
  device_map="auto",
67
  trust_remote_code=True,
68
  ).eval()
 
79
  tokenizer_b = AutoTokenizer.from_pretrained(model_id)
80
  model_b = AutoModelForCausalLM.from_pretrained(
81
  model_id,
82
+ torch_dtype=torch.bfloat16,
83
  device_map="auto",
84
  trust_remote_code=True,
85
  ).eval()