Added MC model to UI and removed some unnecessary code

app.py

@@ -12,10 +12,11 @@ from scipy.special import softmax
 import language_tool_python
 import nltk
 import torch
+import numpy as np
 from transformers import GPT2LMHeadModel, GPT2TokenizerFast
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 
-from utils import remove_special_characters
+from utils import remove_special_characters, split_text_allow_complete_sentences_nltk
 from google_search import google_search, months, domain_list, build_date
 from humanize import paraphrase_text, device
 from ai_generate import generate
@@ -196,65 +197,64 @@ ai_check_options = [
 ]
 
 
-def ai_generated_test_sapling(text: str) -> Dict:
-    response = requests.post(
-        "https://api.sapling.ai/api/v1/aidetect", json={"key": "60L9BPSVPIIOEZM0CD1DQWRBPJIUR7SB", "text": f"{text}"}
-    )
-    return {"AI": response.json()["score"], "HUMAN": 1 - response.json()["score"]}
-
-
-class GPT2PPL:
-    def __init__(self):
-        self.device = device
-        self.model = to_device(GPT2LMHeadModel.from_pretrained("gpt2"))
-        self.tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
+MC_TOKEN_SIZE = 256
+TEXT_MC_MODEL_PATH = "polygraf-ai/mc-model"
+MC_LABEL_MAP = ["OpenAI GPT", "Mistral", "CLAUDE", "Gemini", "Grammar Enhancer"]
+text_mc_tokenizer = AutoTokenizer.from_pretrained(TEXT_MC_MODEL_PATH)
+text_mc_model = AutoModelForSequenceClassification.from_pretrained(TEXT_MC_MODEL_PATH).to(device)
 
-    def __call__(self, text):
-        encodings = self.tokenizer(text, return_tensors="pt")
-        encodings = {k: v.to(self.device) for k, v in encodings.items()}
-        max_length = self.model.config.n_positions
-        stride = 512
-        seq_len = encodings.input_ids.size(1)
 
-        nlls = []
-        for i in range(0, seq_len, stride):
-            begin_loc = max(i + stride - max_length, 0)
-            end_loc = min(i + stride, seq_len)
-            trg_len = end_loc - i
-            input_ids = encodings.input_ids[:, begin_loc:end_loc].to(self.device)
-            target_ids = input_ids.clone()
-            target_ids[:, :-trg_len] = -100
-
-            with torch.no_grad():
-                outputs = self.model(input_ids, labels=target_ids)
-                neg_log_likelihood = outputs.loss * trg_len
+def predict_mc(text):
+    with torch.no_grad():
+        text_mc_model.eval()
+        tokens = text_mc_tokenizer(
+            text,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+            max_length=MC_TOKEN_SIZE,
+        ).to(device)
+        output = text_mc_model(**tokens)
+        output_norm = softmax(output.logits.detach().cpu().numpy(), 1)[0]
+        return output_norm
 
-            nlls.append(neg_log_likelihood)
 
-        ppl = torch.exp(torch.stack(nlls).sum() / end_loc)
-        return {"AI": float(ppl), "HUMAN": 1 - float(ppl)}
+def predict_mc_scores(input, bc_score):
+    mc_scores = []
+    segments_mc = split_text_allow_complete_sentences_nltk(input, type_det="mc", tokenizer=text_mc_tokenizer)
+    samples_len_mc = len(split_text_allow_complete_sentences_nltk(input, type_det="mc", tokenizer=text_mc_tokenizer))
+    for i in range(samples_len_mc):
+        cleaned_text_mc = remove_special_characters(segments_mc[i])
+        mc_score = predict_mc(cleaned_text_mc)
+        mc_scores.append(mc_score)
+    mc_scores_array = np.array(mc_scores)
+    average_mc_scores = np.mean(mc_scores_array, axis=0)
+    mc_score_list = average_mc_scores.tolist()
+    mc_score = {}
+    for score, label in zip(mc_score_list, MC_LABEL_MAP):
+        mc_score[label.upper()] = score
 
+    sum_prob = 1 - bc_score["HUMAN"]
+    for key, value in mc_score.items():
+        mc_score[key] = value * sum_prob
+    print("MC Score:", mc_score)
+    if sum_prob < 0.01:
+        mc_score = {}
 
-def ai_generated_test_gptzero(text):
-    gptzero_model = GPT2PPL()
-    result = gptzero_model(text)
-    return result, None
+    return mc_score
 
 
 def highlighter_polygraf(text, model="Polygraf AI (Base Model)"):
     body, references = split_text_from_refs(text)
     score, text = detection_polygraf(text=body, model=model)
+    mc_score = predict_mc_scores(body, score)  # mc score
     text = text + references.replace("\n", "<br>")
-    return score, text
+    return score, text, mc_score
 
 
 def ai_check(text: str, option: str):
     if option.startswith("Polygraf AI"):
         return highlighter_polygraf(text, option)
-    elif option == "Sapling AI":
-        return ai_generated_test_sapling(text)
-    elif option == "GPTZero":
-        return ai_generated_test_gptzero(text)
     else:
         return highlighter_polygraf(text, option)
 
@@ -788,6 +788,7 @@ def create_interface():
 
                     with gr.Accordion("AI Detection Results", open=True):
                         ai_check_result = gr.Label(label="AI Check Result")
+                        mc_check_result = gr.Label(label="Creator Check Result")
                         highlighted_text = gr.HTML(label="Sentence Breakdown", visible=False)
                     humanize_btn = gr.Button("Humanize")
                     # humanized_output = gr.Markdown(label="Humanized Article", value="\n\n\n\n", render=True)
@@ -904,7 +905,7 @@ def create_interface():
         ai_check_btn.click(
             fn=ai_check,
             inputs=[output_article, ai_detector_dropdown],
-            outputs=[ai_check_result, highlighted_text],
+            outputs=[ai_check_result, highlighted_text, mc_check_result],
         )
 
         humanize_btn.click(

utils.py

@@ -1,5 +1,7 @@
 import re
 from unidecode import unidecode
+from nltk import sent_tokenize
+
 # from transformers import AutoTokenizer
 # import yaml
 # import fitz
@@ -7,33 +9,39 @@ from unidecode import unidecode
 # from bs4 import BeautifulSoup
 # from collections import defaultdict
 
+
 def remove_accents(input_str):
     text_no_accents = unidecode(input_str)
     return text_no_accents
 
+
 def remove_special_characters(text):
-    text = re.sub(r'https?://\S+|www\.\S+', '', text)
-    emoji_pattern = re.compile("["  
-        u"\U0001F600-\U0001F64F"  # emoticons
-        u"\U0001F300-\U0001F5FF"  # symbols & pictographs
-        u"\U0001F680-\U0001F6FF"  # transport & map symbols
-        u"\U0001F700-\U0001F77F"  # alchemical symbols
-        u"\U0001F780-\U0001F7FF"  # Geometric Shapes Extended
-        u"\U0001F800-\U0001F8FF"  # Supplemental Arrows-C
-        u"\U0001F900-\U0001F9FF"  # Supplemental Symbols and Pictographs
-        u"\U0001FA00-\U0001FA6F"  # Chess Symbols
-        u"\U0001FA70-\U0001FAFF"  # Symbols and Pictographs Extended-A
-        u"\U00002702-\U000027B0"  # Dingbats
-        u"\U000024C2-\U0001F251" 
-        "]+", flags=re.UNICODE)
-    text = emoji_pattern.sub('', text) 
-    text = re.sub(r'#\w+', '', text)  
-    text = re.sub(r'[^\w\s\d.,!?\'"()-;]', '', text) 
-    text = re.sub(r'\s+([.,!?;])', r'\1', text)
-    text = re.sub(r'([.,!?;])(\S)', r'\1 \2', text)
-    text = re.sub(r'\s+', ' ', text).strip()
+    text = re.sub(r"https?://\S+|www\.\S+", "", text)
+    emoji_pattern = re.compile(
+        "["
+        "\U0001F600-\U0001F64F"  # emoticons
+        "\U0001F300-\U0001F5FF"  # symbols & pictographs
+        "\U0001F680-\U0001F6FF"  # transport & map symbols
+        "\U0001F700-\U0001F77F"  # alchemical symbols
+        "\U0001F780-\U0001F7FF"  # Geometric Shapes Extended
+        "\U0001F800-\U0001F8FF"  # Supplemental Arrows-C
+        "\U0001F900-\U0001F9FF"  # Supplemental Symbols and Pictographs
+        "\U0001FA00-\U0001FA6F"  # Chess Symbols
+        "\U0001FA70-\U0001FAFF"  # Symbols and Pictographs Extended-A
+        "\U00002702-\U000027B0"  # Dingbats
+        "\U000024C2-\U0001F251"
+        "]+",
+        flags=re.UNICODE,
+    )
+    text = emoji_pattern.sub("", text)
+    text = re.sub(r"#\w+", "", text)
+    text = re.sub(r'[^\w\s\d.,!?\'"()-;]', "", text)
+    text = re.sub(r"\s+([.,!?;])", r"\1", text)
+    text = re.sub(r"([.,!?;])(\S)", r"\1 \2", text)
+    text = re.sub(r"\s+", " ", text).strip()
     return text
 
+
 def remove_special_characters_2(text):
     pattern = r"[^a-zA-Z0-9 ]+"
     text = re.sub(pattern, "", text)
@@ -41,5 +49,56 @@ def remove_special_characters_2(text):
 
 
 def split_into_sentences(text):
-    sentences = re.split(r'(?<=[.!?]) +', text)
-    return sentences
+    sentences = re.split(r"(?<=[.!?]) +", text)
+    return sentences
+
+
+def get_token_length(tokenizer, sentence):
+    return len(tokenizer.tokenize(sentence))
+
+
+MC_TOKEN_SIZE = 256
+BC_TOKEN_SIZE = 333
+
+
+def split_text_allow_complete_sentences_nltk(text, type_det="bc", tokenizer=None):
+    sentences = sent_tokenize(text)
+    chunks = []
+    current_chunk = []
+    current_length = 0
+    if type_det == "bc":
+        max_tokens = BC_TOKEN_SIZE
+    elif type_det == "mc":
+        max_tokens = MC_TOKEN_SIZE
+    elif type_det == "quillbot":
+        max_tokens = 256
+
+    def add_sentence_to_chunk(sentence):
+        nonlocal current_chunk, current_length
+        sentence_length = get_token_length(tokenizer, sentence)
+        if current_length + sentence_length > max_tokens:
+            chunks.append((current_chunk, current_length))
+            current_chunk = []
+            current_length = 0
+        current_chunk.append(sentence)
+        current_length += sentence_length
+
+    for sentence in sentences:
+        add_sentence_to_chunk(sentence)
+    if current_chunk:
+        chunks.append((current_chunk, current_length))
+    adjusted_chunks = []
+    while chunks:
+        chunk = chunks.pop(0)
+        if len(chunks) > 0 and chunk[1] < max_tokens / 2:
+            next_chunk = chunks.pop(0)
+            combined_length = chunk[1] + next_chunk[1]
+            if combined_length <= max_tokens:
+                adjusted_chunks.append((chunk[0] + next_chunk[0], combined_length))
+            else:
+                adjusted_chunks.append(chunk)
+                chunks.insert(0, next_chunk)
+        else:
+            adjusted_chunks.append(chunk)
+    result_chunks = [" ".join(chunk[0]) for chunk in adjusted_chunks]
+    return result_chunks