0x7o commited on
Commit
ed96116
1 Parent(s): e738d31

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -6,9 +6,9 @@ from transformers import StoppingCriteria, StoppingCriteriaList, TextIteratorStr
6
  from threading import Thread
7
 
8
  # Loading the tokenizer and model from Hugging Face's model hub.
9
- tokenizer = AutoTokenizer.from_pretrained("0x7194633/fialka-13B-v3.1")
10
- model = AutoModelForCausalLM.from_pretrained("0x7194633/fialka-13B-v3.1", load_in_8bit=True)
11
- model.to("cuda:0")
12
 
13
 
14
  # Defining a custom stopping criteria class for the model's text generation.
 
6
  from threading import Thread
7
 
8
  # Loading the tokenizer and model from Hugging Face's model hub.
9
+ if torch.cuda.is_available():
10
+ tokenizer = AutoTokenizer.from_pretrained("0x7194633/fialka-13B-v3.1")
11
+ model = AutoModelForCausalLM.from_pretrained("0x7194633/fialka-13B-v3.1", load_in_8bit=True, device_map="auto")
12
 
13
 
14
  # Defining a custom stopping criteria class for the model's text generation.