Change gen output

app.py

@@ -16,6 +16,7 @@ tokenizer = AutoTokenizer.from_pretrained(peft_model_id)
 # Load the Lora model
 model = PeftModel.from_pretrained(model, peft_model_id)
 
+model.eval()
 
 def gen_entities(text):
     text = f"<SP> text: {text}\n\n entities: "
@@ -23,7 +24,8 @@ def gen_entities(text):
     with torch.cuda.amp.autocast():
         output_tokens = model.generate(**batch, max_new_tokens=256, eos_token_id=50258)
 
-    return tokenizer.decode(output_tokens[0], skip_special_tokens=False)
+    # return tokenizer.decode(output_tokens, skip_special_tokens=False)
+    return tokenizer.batch_decode(output_tokens.detach().cpu().numpy(), skip_special_tokens=True)
 
 
 iface = gr.Interface(fn=gen_entities, inputs="text", outputs="text")