Change gen_entities function

app.py

@@ -18,14 +18,20 @@ model = PeftModel.from_pretrained(model, peft_model_id)
 
 model.eval()
 
-def gen_entities(text):
-    text = f"<SP> text: {text}\n\n entities: "
+
+def gen_entities(in_text):
+    """Does Named Entity Recognition in the given text."""
+    text = f"<SP> text: {in_text}\n\n entities:"
     batch = tokenizer(text, return_tensors="pt")
+    batch["input_ids"] = batch["input_ids"].to("cuda")
     with torch.cuda.amp.autocast():
         output_tokens = model.generate(**batch, max_new_tokens=256, eos_token_id=50258)
 
-    # return tokenizer.decode(output_tokens, skip_special_tokens=False)
-    return tokenizer.batch_decode(output_tokens.detach().cpu().numpy(), skip_special_tokens=True)
+    response = tokenizer.batch_decode(
+        output_tokens.detach().cpu().numpy(), skip_special_tokens=False
+    )[0]
+
+    return response[response.find("entities") : response.find("<EP>")]
 
 
 iface = gr.Interface(fn=gen_entities, inputs="text", outputs="text")