Update app.py

app.py

@@ -29,6 +29,7 @@ tokenizer = AutoTokenizer.from_pretrained("checkpoint_500", trust_remote_code=Tr
 tokenizer.pad_token = tokenizer.eos_token
 
 def inference(prompt, count):
+    count = int(count)
     pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=200) 
     result = pipe(f"### Human: {prompt}",max_new_tokens=count)
     out_text = result[0]['generated_text']