Spaces:

joaogante
/

assisted_generation_demo

Running on Zero

joaogante HF Staff commited on Mar 4

Commit

b1b72e8

verified ·

1 Parent(s): 817dcba

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -8,7 +8,7 @@ import time
 model_id = "facebook/opt-6.7b"
 assistant_id = "facebook/opt-125m"
-model = AutoModelForCausalLM.from_pretrained(model_id, load_in_8bit=True, device_map="auto")
 assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id).to(device=model.device, dtype=torch.bfloat16)
 tokenizer = AutoTokenizer.from_pretrained(model_id)

 model_id = "facebook/opt-6.7b"
 assistant_id = "facebook/opt-125m"
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
 assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id).to(device=model.device, dtype=torch.bfloat16)
 tokenizer = AutoTokenizer.from_pretrained(model_id)