Rajesh3338 commited on
Commit
2e56b20
·
verified ·
1 Parent(s): 637932e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -30,14 +30,14 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
30
  model = AutoModelForCausalLM.from_pretrained(
31
  model_name,
32
  device_map=device,
33
- torch_dtype=torch.float32 if device == "cuda" else torch.float32
34
  )
35
  #model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype="auto")
36
  qa_pipeline = pipeline(
37
  "text-generation",
38
  model=model,
39
  tokenizer=tokenizer,
40
- max_new_tokens=500,
41
  pad_token_id=tokenizer.eos_token_id
42
  )
43
 
 
30
  model = AutoModelForCausalLM.from_pretrained(
31
  model_name,
32
  device_map=device,
33
+ torch_dtype=torch.float16 if device == "cuda" else torch.float32
34
  )
35
  #model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype="auto")
36
  qa_pipeline = pipeline(
37
  "text-generation",
38
  model=model,
39
  tokenizer=tokenizer,
40
+ max_new_tokens=1500,
41
  pad_token_id=tokenizer.eos_token_id
42
  )
43