Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -43,12 +43,12 @@ print(f"Starting to load the model to memory")
|
|
43 |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
44 |
|
45 |
m = AutoModelForCausalLM.from_pretrained(
|
46 |
-
"google/gemma-2b-it", torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, trust_remote_code=True
|
47 |
embedding_func=m.get_input_embeddings()
|
48 |
embedding_func.weight.requires_grad=False
|
49 |
m = m.to(device)
|
50 |
|
51 |
-
tok = AutoTokenizer.from_pretrained("google/gemma-2b-it", trust_remote_code=True
|
52 |
tok.padding_side = "left"
|
53 |
tok.pad_token_id = tok.eos_token_id
|
54 |
# using CUDA for an optimal experience
|
|
|
43 |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
44 |
|
45 |
m = AutoModelForCausalLM.from_pretrained(
|
46 |
+
"google/gemma-2b-it", torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, trust_remote_code=True)
|
47 |
embedding_func=m.get_input_embeddings()
|
48 |
embedding_func.weight.requires_grad=False
|
49 |
m = m.to(device)
|
50 |
|
51 |
+
tok = AutoTokenizer.from_pretrained("google/gemma-2b-it", trust_remote_code=True)
|
52 |
tok.padding_side = "left"
|
53 |
tok.pad_token_id = tok.eos_token_id
|
54 |
# using CUDA for an optimal experience
|