tungdop2 commited on
Commit
b35fa97
1 Parent(s): 486a2da

fix docker

Browse files
Files changed (1) hide show
  1. model.py +3 -3
model.py CHANGED
@@ -22,12 +22,12 @@ class ChallengePromptGenerator:
22
  self.tokenizer = AutoTokenizer.from_pretrained(model_local_dir)
23
 
24
  def infer_prompt(
25
- self,
26
  prompts,
27
  max_generation_length=77,
28
  beam_size=1,
29
  sampling_temperature=0.9,
30
- sampling_topk=1,
31
  sampling_topp=1
32
  ):
33
  # Add bos
@@ -62,4 +62,4 @@ class ChallengePromptGenerator:
62
  if out[-1] != ".":
63
  out = ".".join(out.split(".")[:-1]) + "."
64
  outputs.append(out)
65
- return decoded_outputs
 
22
  self.tokenizer = AutoTokenizer.from_pretrained(model_local_dir)
23
 
24
  def infer_prompt(
25
+ self,
26
  prompts,
27
  max_generation_length=77,
28
  beam_size=1,
29
  sampling_temperature=0.9,
30
+ sampling_topk=100,
31
  sampling_topp=1
32
  ):
33
  # Add bos
 
62
  if out[-1] != ".":
63
  out = ".".join(out.split(".")[:-1]) + "."
64
  outputs.append(out)
65
+ return outputs