IMJONEZZ commited on
Commit
0a39f44
·
verified ·
1 Parent(s): 9d139d0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -1,10 +1,10 @@
1
  import gradio as gr
2
- import torch
3
  from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
4
  from threading import Thread
 
5
 
6
  tokenizer = AutoTokenizer.from_pretrained("./model/")
7
- model = AutoModelForCausalLM.from_pretrained("./model/", torch_dtype=torch.float16)
8
  model = model.to('cuda:0')
9
 
10
  class StopOnTokens(StoppingCriteria):
 
1
  import gradio as gr
 
2
  from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
3
  from threading import Thread
4
+ import bitsandbytes
5
 
6
  tokenizer = AutoTokenizer.from_pretrained("./model/")
7
+ model = AutoModelForCausalLM.from_pretrained("./model/", device_map="auto", load_in_4bit=True)
8
  model = model.to('cuda:0')
9
 
10
  class StopOnTokens(StoppingCriteria):