edithram23 commited on
Commit
31a390f
·
verified ·
1 Parent(s): e9a0520

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +1 -1
main.py CHANGED
@@ -12,7 +12,7 @@ def mask_generation(text):
12
  import re
13
  inputs = ["Mask Generation: " + text]
14
  inputs = tokenizer(inputs, max_length=128, truncation=True, return_tensors="pt")
15
- output = model.generate(**inputs, num_beams=8, do_sample=True, max_length=len(text))
16
  decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
17
  predicted_title = decoded_output.strip()
18
  pattern = r'\[.*?\]'
 
12
  import re
13
  inputs = ["Mask Generation: " + text]
14
  inputs = tokenizer(inputs, max_length=128, truncation=True, return_tensors="pt")
15
+ output = model.generate(**inputs, num_beams=8, do_sample=True, max_length=64)
16
  decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
17
  predicted_title = decoded_output.strip()
18
  pattern = r'\[.*?\]'