from transformers import pipeline
import gradio as gr

examples = [
"""
Notice of Representation 
 Tachyon & Park PLLC 
 1618 25th Ave 
 Spokane, Washington(WA), 99208 
 Direct Insurance Company 
 5555 Dakota St. 
 Athens, GA 23001 
 Re: Estate of Bryan Terrell 
 Policy Number: 117213657 
 Date of death: 6/8/2021 
 To Whom It May Concern: 
 I have been retained by Saskia Mcgee to handle the abovenamed estate. My understanding is 
 that they had a life insurance policy with your company. If this is correct, please send a letter to 
 my office indicating you have received our letter of representation. Additionally, please do not 
 contact our client going forward. 
 We request that you deliver posthaste, the full policy amount of $400,000. If you are aware of 
 any additional policies in force, please provide us with that information. Additionally, if there 
 are any exclusions or liens on the policy, we request that information as well. 
 If you have any questions, please contact my office. 
 Sincerely, 
 J Rock, esq. 
""",
    """
 Notice of Representation 
 Budget Mutual Insurance Company 
 9876 Infinity Ave 
 Springfield, MI 65541 
 Colin & Bryier PLLC 
 9514 8th Ave S 
 Auburn, Washington(WA), 98002 
 Our Client: Aysha Gilmore 
 Date of death: 7/8/2021 
 To Whom It May Concern, 
 I have been retained by Aysha Gilmore to handle the estate of Kyron Marks. My understanding 
 is that they had a life insurance policy (#193635138) with your company. If this is correct, 
 please send a letter to my office indicating you have received our letter of representation. 
 Additionally, please do not contact our client going forward. 
 We are requesting that you forward the full policy amount of $25,000. Please forward an 
 acknowledgement of our demand and please forward the umbrella policy information if one is 
 applicable. Please send my secretary any information regarding liens on his policy. 
 Please contact my office if you have any questions. 
 Sincerely, 
 Angela Berry, Attorney""",
"""
Notice of Representation 
 Number One Insurance Company 
 1234 Gateway Dr 
 Chicago, IL 15002 
 Quiroga PLLC 
 9668 Rainier Ave S 
 Kent, Washington(WA), 98031 
 Re: Estate of Sana Keith 
 Policy number: 462204232 
 Our client: Oliver Davis 
 Date of death: 2/14/2020 
 To Whom It May Concern, 
 I have been retained by Oliver Davis to handle the estate of Sana Keith. My understanding is 
 that they had a life insurance policy with your company. If this is correct, please send a letter to 
 my office indicating you have received our letter of representation. Additionally, please do not 
 contact our client going forward. 
 Our understanding is that the policy was for the amount of $60,000. If that is correct, please 
 forward that amount to our office. If there are any forms that need to be completed, please 
 forward those as well. If you are aware of any additional policies that are in force, send 
 information about those policies to our office. 
 If you have any questions, please contact my office. 
 Sincerely, 
 John D Locke, Esq"""
]
# Replace this with your own checkpoint
model_checkpoint = "aimlnerd/bert-finetuned-legalentity-ner"
ner_pipeline = pipeline(
    "token-classification", model=model_checkpoint, aggregation_strategy="simple"
)

def ner(text):
    output = ner_pipeline(text)
    return {"text": text, "entities": output}

demo = gr.Interface(ner,
             gr.Textbox(placeholder="Enter text here..."), 
             gr.HighlightedText(),
             examples=examples)

demo.launch(share=True)
print()