Spaces:
Paused
Paused
Commit
·
6de173e
1
Parent(s):
5d04704
Update app_v4.py
Browse files
app_v4.py
CHANGED
|
@@ -33,7 +33,7 @@ if device == "cuda:0":
|
|
| 33 |
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=False)
|
| 34 |
|
| 35 |
# Attempt to load the model, catch any OOM errors
|
| 36 |
-
@st.
|
| 37 |
def load_gptq_model():
|
| 38 |
model = AutoGPTQForCausalLM.from_quantized(
|
| 39 |
pretrained_model_dir,
|
|
@@ -45,10 +45,8 @@ def load_gptq_model():
|
|
| 45 |
return model
|
| 46 |
|
| 47 |
# Attempt to load the model, catch any OOM errors
|
| 48 |
-
model_loaded = False
|
| 49 |
try:
|
| 50 |
model = load_gptq_model()
|
| 51 |
-
model_loaded = True
|
| 52 |
except RuntimeError as e:
|
| 53 |
if 'CUDA out of memory' in str(e):
|
| 54 |
st.error("CUDA out of memory while loading the model. Try reducing the model size or restarting the app.")
|
|
|
|
| 33 |
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=False)
|
| 34 |
|
| 35 |
# Attempt to load the model, catch any OOM errors
|
| 36 |
+
@st.cache_resource
|
| 37 |
def load_gptq_model():
|
| 38 |
model = AutoGPTQForCausalLM.from_quantized(
|
| 39 |
pretrained_model_dir,
|
|
|
|
| 45 |
return model
|
| 46 |
|
| 47 |
# Attempt to load the model, catch any OOM errors
|
|
|
|
| 48 |
try:
|
| 49 |
model = load_gptq_model()
|
|
|
|
| 50 |
except RuntimeError as e:
|
| 51 |
if 'CUDA out of memory' in str(e):
|
| 52 |
st.error("CUDA out of memory while loading the model. Try reducing the model size or restarting the app.")
|