Spaces:
Running
on
Zero
Running
on
Zero
AstraBert
commited on
Commit
•
5239972
1
Parent(s):
d107425
bloom
Browse files- load_model.py +3 -3
load_model.py
CHANGED
@@ -1,6 +1,6 @@
|
|
1 |
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
|
2 |
|
3 |
-
model = AutoModelForCausalLM.from_pretrained("
|
4 |
-
tokenizer = AutoTokenizer.from_pretrained("
|
5 |
|
6 |
-
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=2048)
|
|
|
1 |
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
|
2 |
|
3 |
+
model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-1b7")
|
4 |
+
tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-1b7")
|
5 |
|
6 |
+
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=2048, repetition_penalty=1.2)
|