src/texts/MAGE/app.py

from transformers import pipeline
from difflib import Differ
from transformers import AutoModelForSequenceClassification,AutoTokenizer
from deployment import preprocess, detect
import gradio as gr

ner_pipeline = pipeline("ner")


def ner(text):
    output = ner_pipeline(text)
    output = [
        {'entity': 'I-LOC', 'score': 0.9995369, 'index': 2, 'word': 'Chicago', 'start': 5, 'end': 12},
        {'entity': 'I-PER', 'score': 0.99527764, 'index': 8, 'word': 'Joe', 'start': 38, 'end': 41}
        ]
    print(output)
    return {"text": text, "entities": output}

def diff_texts(text1, text2):
    d = Differ()
    return [
        (token[2:], token[0] if token[0] != " " else None)
        for token in d.compare(text1, text2)
    ]
    
out = diff_texts(
    "The quick brown fox jumped over the lazy dogs.",
    "The fast brown fox jumps over lazy dogs.")
print(out)


def separate_characters_with_mask(text, mask):
  """Separates characters in a string and pairs them with a mask sign.

  Args:
    text: The input string.

  Returns:
    A list of tuples, where each tuple contains a character and a mask.
  """

  return [(char, mask) for char in text]


def detect_ai_text(text):
    text = preprocess(text)
    result = detect(text,tokenizer,model,device)
    print(result)
    output = separate_characters_with_mask(text, result)
    return output

# init
device = 'cpu' # use 'cuda:0' if GPU is available
# model_dir = "nealcly/detection-longformer" # model in our paper
model_dir = "yaful/MAGE" # model in the online demo
tokenizer = AutoTokenizer.from_pretrained(model_dir)
model = AutoModelForSequenceClassification.from_pretrained(model_dir).to(device)
examples = ["Apple's new credit card will begin a preview roll out today and will become available to all iPhone owners in the US later this month. A random selection of people will be allowed to go through the application process, which involves entering personal details which are sent to Goldman Sachs and TransUnion. Applications are approved or declined in less than a minute. The Apple Card is meant to be broadly accessible to every iPhone user, so the approval requirements will not be as strict as other credit cards. Once the application has been approved, users will be able to use the card immediately from the Apple Wallet app. The physical titanium card can be requested during setup for free, and it can be activated with NFC once it arrives."]

demo = gr.Interface(detect_ai_text,
             gr.Textbox(
                 label="input text",
                 placeholder="Enter text here...",
                 lines=5,
                 ),
             gr.HighlightedText(
                 label="AI-text detection",
                 combine_adjacent=True,
                 show_legend=True,
                 color_map={"machine-generated": "red", "human-written": "green"}
                 ),
             examples=examples)

demo.launch(share=True)