Spaces:
Runtime error
Runtime error
stefiane Zhang (ๅผ ๅฎถๅ)-ๆตชๆฝฎไฟกๆฏ
commited on
Commit
ยท
9ab4de0
1
Parent(s):
82b6228
Add application file
Browse files- app.py +3 -1
- requirements.txt +2 -1
app.py
CHANGED
@@ -15,7 +15,9 @@ tokenizer.add_tokens(['<sep>', '<pad>', '<mask>', '<predict>', '<FIM_SUFFIX>', '
|
|
15 |
|
16 |
print("Creat model...")
|
17 |
model = AutoModelForCausalLM.from_pretrained('IEITYuan/Yuan2-2B-hf', device_map='auto', torch_dtype=torch.bfloat16, trust_remote_code=True)
|
18 |
-
|
|
|
|
|
19 |
|
20 |
# Defining a custom stopping criteria class for the model's text generation.
|
21 |
class StopOnTokens(StoppingCriteria):
|
|
|
15 |
|
16 |
print("Creat model...")
|
17 |
model = AutoModelForCausalLM.from_pretrained('IEITYuan/Yuan2-2B-hf', device_map='auto', torch_dtype=torch.bfloat16, trust_remote_code=True)
|
18 |
+
# using CUDA for an optimal experience
|
19 |
+
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
20 |
+
model = model.to(device)
|
21 |
|
22 |
# Defining a custom stopping criteria class for the model's text generation.
|
23 |
class StopOnTokens(StoppingCriteria):
|
requirements.txt
CHANGED
@@ -1,3 +1,4 @@
|
|
1 |
torch>=2.0
|
2 |
transformers>=4.35.0
|
3 |
-
gradio>=4.13.0
|
|
|
|
1 |
torch>=2.0
|
2 |
transformers>=4.35.0
|
3 |
+
gradio>=4.13.0
|
4 |
+
sentencepiece
|