import gradio as gr
import requests


def extract_pii(text):
    res = requests.post(
        "https://6b044.cognitiveservices.azure.com/language/:analyze-text?api-version=2022-05-01",
        headers={
            "Ocp-Apim-Subscription-Key": "941b608ce2634805b169eca63f51b856",
            "Content-Type": "application/json",
        },
        json={
            "kind": "PiiEntityRecognition",
            "parameters": {
                "modelVersion": "latest",
            },
            "analysisInput":{
                "documents":[
                    {
                        "id":"1",
                        "language": "en",
                        "text": text,
                    },
                ],
            },
        },
    )
    pii_entities = res.json()["results"]["documents"][0]["entities"]
    return list(map(lambda row: [row["text"], row["category"], row["confidenceScore"]], pii_entities))

demo = gr.Interface(extract_pii, inputs=gr.Textbox(label="Text for PII extraction"), outputs=gr.DataFrame(headers=["Text", "Category", "Confidence Score"]))

if __name__ == "__main__":
    demo.launch()