Update app.py
Browse files
app.py
CHANGED
@@ -29,6 +29,7 @@ tokenizer = AutoTokenizer.from_pretrained("checkpoint_500", trust_remote_code=Tr
|
|
29 |
tokenizer.pad_token = tokenizer.eos_token
|
30 |
|
31 |
def inference(prompt, count):
|
|
|
32 |
pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=200)
|
33 |
result = pipe(f"### Human: {prompt}",max_new_tokens=count)
|
34 |
out_text = result[0]['generated_text']
|
|
|
29 |
tokenizer.pad_token = tokenizer.eos_token
|
30 |
|
31 |
def inference(prompt, count):
|
32 |
+
count = int(count)
|
33 |
pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=200)
|
34 |
result = pipe(f"### Human: {prompt}",max_new_tokens=count)
|
35 |
out_text = result[0]['generated_text']
|