chansung commited on
Commit
c582d61
1 Parent(s): 953dd6b

Update gen.py

Browse files
Files changed (1) hide show
  1. gen.py +8 -2
gen.py CHANGED
@@ -81,11 +81,17 @@ def get_pretrained_models(
81
 
82
  def get_output(
83
  generator: LLaMA,
84
- prompt: str,
 
85
  temperature: float = 0.8,
86
  top_p: float = 0.95):
87
 
88
  prompts = [prompt]
89
- results = generator.generate(prompts, max_gen_len=256, temperature=temperature, top_p=top_p)
 
 
 
 
 
90
 
91
  return results
 
81
 
82
  def get_output(
83
  generator: LLaMA,
84
+ prompt: str,
85
+ max_gen_len: int = 256,
86
  temperature: float = 0.8,
87
  top_p: float = 0.95):
88
 
89
  prompts = [prompt]
90
+ results = generator.generate(
91
+ prompts,
92
+ max_gen_len=max_gen_len,
93
+ temperature=temperature,
94
+ top_p=top_p
95
+ )
96
 
97
  return results