YANGSongsong commited on
Commit
07b3063
1 Parent(s): 865adee

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -9,7 +9,8 @@ TOKENIZER_PATH = os.environ.get("TOKENIZER_PATH", MODEL_PATH)
9
 
10
  tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH, trust_remote_code=True)
11
 
12
- model = AutoModel.from_pretrained(MODEL_PATH, trust_remote_code=True).float()
 
13
 
14
  class StopOnTokens(StoppingCriteria):
15
  def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
 
9
 
10
  tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH, trust_remote_code=True)
11
 
12
+ # model = AutoModel.from_pretrained(MODEL_PATH, trust_remote_code=True).float()
13
+ model = AutoModel.from_pretrained(MODEL_PATH, trust_remote_code=True, device_map="auto").eval()
14
 
15
  class StopOnTokens(StoppingCriteria):
16
  def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool: