Spaces:

walledai
/

walledai-walledguard-c

Sleeping

RishabhBhardwaj commited on Jul 24, 2024

Commit

7998472

1 Parent(s): 468731a

load 4 bit

Files changed (1) hide show

app.py CHANGED Viewed

@@ -18,7 +18,7 @@ Answer: [/INST]
 def load_model():
     model_name = "meta-llama/Llama-Guard-3-8B"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
     return tokenizer, model
 # Function to load image from URL

 def load_model():
     model_name = "meta-llama/Llama-Guard-3-8B"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name, load_in_4bit=True)
     return tokenizer, model
 # Function to load image from URL