kouki321 commited on
Commit
9dfced6
·
verified ·
1 Parent(s): 2fb73b9

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -88,8 +88,9 @@ def load_model_and_tokenizer(doc_text_count):
88
  # Load the pre-trained model with quantization
89
  model = AutoModelForCausalLM.from_pretrained(
90
  model_name,
91
- trust_remote_code=True,
92
- model_max_length=1.3*round(doc_text_count * 0.3 + 1)
 
93
  ,token=HF_TOKEN
94
  )
95
 
 
88
  # Load the pre-trained model with quantization
89
  model = AutoModelForCausalLM.from_pretrained(
90
  model_name,
91
+ torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
92
+ device_map="auto",
93
+ trust_remote_code=True
94
  ,token=HF_TOKEN
95
  )
96