leonardlin commited on
Commit
0b5b812
·
verified ·
1 Parent(s): 74b5f00

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -2,6 +2,7 @@ import os
2
  from collections.abc import Iterator
3
  from threading import Thread
4
 
 
5
  import gradio as gr
6
  import spaces
7
  import torch
@@ -25,7 +26,7 @@ if not torch.cuda.is_available():
25
 
26
  if torch.cuda.is_available():
27
  model_id = "shisa-ai/shisa-v2-unphi-14b-W8A8-INT8"
28
- model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto", load_in_8bit=True)
29
  tokenizer = AutoTokenizer.from_pretrained(model_id)
30
  tokenizer.use_default_system_prompt = False
31
 
 
2
  from collections.abc import Iterator
3
  from threading import Thread
4
 
5
+ import compressed_tensors
6
  import gradio as gr
7
  import spaces
8
  import torch
 
26
 
27
  if torch.cuda.is_available():
28
  model_id = "shisa-ai/shisa-v2-unphi-14b-W8A8-INT8"
29
+ model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto", load_in_8bit=True, trust_remote_code=True)
30
  tokenizer = AutoTokenizer.from_pretrained(model_id)
31
  tokenizer.use_default_system_prompt = False
32