import gradio as gr
from transformers import pipeline
import spacy
import spacy.cli
spacy.cli.download("en_core_web_lg")
nlp = spacy.load("en_core_web_lg")


model_checkpoint ="Umesh/bert-pulc"
token_classifier = pipeline("text-classification", model=model_checkpoint, aggregation_strategy="simple")


def clean_text(text):
    '''
    Cleans the input text
    '''
    text = re.sub(r'\\n', ' ', text)
    text = text.encode('ascii', 'ignore').decode()

    text = text.strip()      
    
    text=re.sub(r'http\S+', '', text)

    text = re.sub("[^,.:1-9A-Za-z]+", " ", text)

    text = re.sub(" +", " ", text)

    return text

def main(everyline):
    

    doc = nlp(everyline) #identifies named entities
    #print('Entities', [(ent.text, ent.label_) for ent in doc.ents])
    #print('Tokens', [(t.text, t.ent_type_, t.ent_iob) for t in doc])
    #print("------------------------------")
    finalString = ""
    for t in doc:
        
        # skipping inside labels
        if int(t.ent_iob) == 1:
            continue
        
        
        # this takes care of the first piece of string in a text
        # Also we identify the beginning label which is label 3 in the following if else statement.
        if finalString == "":
            if str(t.ent_type_) == "PERSON":
                finalString = "<person>"
            elif str(t.ent_type_) == 'GPE':
                finalString = "<loc>"
            elif str(t.ent_type_) == 'LOC':
                finalString = "<loc>"
            elif str(t.ent_type_) == 'NORP':
                finalString = "<norp>"
            elif str(t.ent_type_) == 'FAC':
                finalString = "<fac>"
            elif str(t.ent_type_) == 'ORG':
                finalString = "<org>"
            else:
                finalString = str(t.text)
        elif int(t.ent_iob) == 3:
            
            if str(t.ent_type_) == "PERSON":
                finalString += " <person>"
            elif str(t.ent_type_) == 'GPE':
                finalString += " <loc>"
            elif str(t.ent_type_) == 'LOC':
                finalString += " <loc>"
            elif str(t.ent_type_) == 'NORP':
                finalString += " <norp>"
            elif str(t.ent_type_) == 'FAC':
                finalString += " <fac>"
            elif str(t.ent_type_) == 'ORG':
                finalString += " <org>"
            else:
                finalString += " " + str(t.text)
        else:
            finalString += " " + str(t.text)
    
    return finalString



def inference(text):
    output = token_classifier(text)
    return {p["label"]: p["score"] for p in output}




demo = gr.Interface(inference,
             gr.Textbox(placeholder="Enter sentence here..."), 
             gr.HighlightedText(),
             allow_flagging = False,
             title="Text Classification-PULF")

demo.launch()