Pmal commited on
Commit
f0f445c
·
verified ·
1 Parent(s): fd165fd

reduced size

Browse files
Files changed (1) hide show
  1. app.py +5 -2
app.py CHANGED
@@ -15,9 +15,12 @@ from transformers import pipeline
15
  from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
16
 
17
  def load_model():
18
- tokenizer = AutoTokenizer.from_pretrained("AI-MO/NuminaMath-72B-CoT")
 
 
19
  model = AutoModelForCausalLM.from_pretrained(
20
- "AI-MO/NuminaMath-72B-CoT",
 
21
  device_map="auto", # Automatically map to available GPU
22
  offload_folder="offload" # Offload unused parts to disk
23
  )
 
15
  from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
16
 
17
  def load_model():
18
+ # tokenizer = AutoTokenizer.from_pretrained("AI-MO/NuminaMath-72B-CoT")
19
+ tokenizer = AutoTokenizer.from_pretrained("AI-MO/NuminaMath-7B-CoT")
20
+
21
  model = AutoModelForCausalLM.from_pretrained(
22
+ # "AI-MO/NuminaMath-72B-CoT",
23
+ "AI-MO/NuminaMath-7B-CoT",
24
  device_map="auto", # Automatically map to available GPU
25
  offload_folder="offload" # Offload unused parts to disk
26
  )