Update gen.py

gen.py

@@ -81,11 +81,17 @@ def get_pretrained_models(
 
 def get_output(
     generator: LLaMA,
-    prompt: str, 
+    prompt: str,
+    max_gen_len: int = 256,
     temperature: float = 0.8, 
     top_p: float = 0.95):
     
     prompts = [prompt]
-    results = generator.generate(prompts, max_gen_len=256, temperature=temperature, top_p=top_p)
+    results = generator.generate(
+        prompts, 
+        max_gen_len=max_gen_len, 
+        temperature=temperature, 
+        top_p=top_p
+    )
 
     return results