Spaces:
Sleeping
Sleeping
0.14 explictly set GPU and dtype
Browse files
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=
|
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=
|
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()
|