Merge branch 'main' into demo

app.py

@@ -17,6 +17,11 @@ import fitz
 from transformers import GPT2LMHeadModel, GPT2TokenizerFast
 import nltk, spacy, subprocess, torch
 import plotly.graph_objects as go
+import nltk
+
+nltk.download('punkt')
+tokenizer = AutoTokenizer.from_pretrained('google-bert/bert-base-uncased')
+
 from writing_analysis import (
     normalize,
     preprocess_text1,
@@ -175,11 +180,11 @@ AI DETECTION SECTION
 """
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-text_bc_model_path = "polygraf-ai/ai-text-bc-bert-2-7m"
+text_bc_model_path = "polygraf-ai/v3-bert-3-2m-trun-bc"
 text_bc_tokenizer = AutoTokenizer.from_pretrained(text_bc_model_path)
 text_bc_model = AutoModelForSequenceClassification.from_pretrained(text_bc_model_path).to(device)
 
-text_mc_model_path = "polygraf-ai/ai-text-mc-v5-lighter-spec"
+text_mc_model_path = "polygraf-ai/text-detect-mc-bert-base-uncased-v1-bert-429k"
 text_mc_tokenizer = AutoTokenizer.from_pretrained(text_mc_model_path)
 text_mc_model = AutoModelForSequenceClassification.from_pretrained(text_mc_model_path).to(device)
 
@@ -190,34 +195,105 @@ def remove_special_characters(text):
 def update_character_count(text):
     return f"{len(text)} characters"
 
+
+def split_text_allow_complete_sentences_nltk(text, max_length=256, tolerance=10, min_last_segment_length=120):
+    sentences = nltk.sent_tokenize(text)
+    segments = []
+    current_segment = []
+    current_length = 0  
+    
+    for sentence in sentences:
+        tokens = tokenizer.tokenize(sentence)
+        sentence_length = len(tokens)
+        
+        if current_length + sentence_length <= max_length + tolerance - 2: 
+            current_segment.append(sentence)
+            current_length += sentence_length
+        else:
+            if current_segment:
+                encoded_segment = tokenizer.encode(' '.join(current_segment), add_special_tokens=True, max_length=max_length+tolerance, truncation=True)
+                segments.append((current_segment, len(encoded_segment)))
+            current_segment = [sentence]
+            current_length = sentence_length
+    
+    if current_segment:
+        encoded_segment = tokenizer.encode(' '.join(current_segment), add_special_tokens=True, max_length=max_length+tolerance, truncation=True)
+        segments.append((current_segment, len(encoded_segment)))
+
+    final_segments = []
+    for i, (seg, length) in enumerate(segments):
+        if i == len(segments) - 1:  
+            if length < min_last_segment_length and len(final_segments) > 0:
+                prev_seg, prev_length = final_segments[-1]
+                combined_encoded = tokenizer.encode(' '.join(prev_seg + seg), add_special_tokens=True, max_length=max_length+tolerance, truncation=True)
+                if len(combined_encoded) <= max_length + tolerance:
+                    final_segments[-1] = (prev_seg + seg, len(combined_encoded))
+                else:
+                    final_segments.append((seg, length))
+            else:
+                final_segments.append((seg, length))
+        else:
+            final_segments.append((seg, length))
+
+    decoded_segments = []
+    encoded_segments = []
+    for seg, _ in final_segments:
+        encoded_segment = tokenizer.encode(' '.join(seg), add_special_tokens=True, max_length=max_length+tolerance, truncation=True)
+        decoded_segment = tokenizer.decode(encoded_segment)
+        decoded_segments.append(decoded_segment)
+    return decoded_segments
+
+
 def predict_bc(model, tokenizer, text):
     tokens = tokenizer(
-        text, padding=True, truncation=True, return_tensors="pt"
+        text, padding='max_length', truncation=True, max_length=256, return_tensors="pt"
     ).to(device)["input_ids"]
+    
     output = model(tokens)
     output_norm = softmax(output.logits.detach().cpu().numpy(), 1)[0]
     print("BC Score: ", output_norm)
-    bc_score = {"AI": output_norm[1].item(), "HUMAN": output_norm[0].item()}
-    return bc_score
+    return output_norm
 
 def predict_mc(model, tokenizer, text):
     tokens = tokenizer(
-        text, padding=True, truncation=True, return_tensors="pt"
+        text, padding='max_length', truncation=True, return_tensors="pt", max_length=512
     ).to(device)["input_ids"]
     output = model(tokens)
     output_norm = softmax(output.logits.detach().cpu().numpy(), 1)[0]
     print("MC Score: ", output_norm)
-    mc_score = {}
-    label_map = ["GPT 3.5", "GPT 4", "CLAUDE", "BARD", "LLAMA 2"]
-    for score, label in zip(output_norm, label_map):
-        mc_score[label.upper()] = score.item()
-    return mc_score
+    return output_norm
 
 def ai_generated_test(ai_option, input):
-
-    cleaned_text = remove_special_characters(input)
-    bc_score = predict_bc(text_bc_model, text_bc_tokenizer, cleaned_text)
-    mc_score = predict_mc(text_mc_model, text_mc_tokenizer, cleaned_text)
+    
+    bc_scores = []
+    mc_scores = []
+    samples_len = len(split_text_allow_complete_sentences_nltk(input))
+    segments = split_text_allow_complete_sentences_nltk(input)
+    
+    for i in range(samples_len):
+        cleaned_text = remove_special_characters(segments[i])
+        bc_score = predict_bc(text_bc_model, text_bc_tokenizer,cleaned_text )
+        mc_score = predict_mc(text_mc_model, text_mc_tokenizer, cleaned_text)
+        bc_scores.append(bc_score)
+        mc_scores.append(mc_score)
+
+    bc_scores_array = np.array(bc_scores)
+    mc_scores_array = np.array(mc_scores)
+    average_bc_scores = np.mean(bc_scores_array, axis=0)
+    average_mc_scores = np.mean(mc_scores_array, axis=0)
+    bc_score_list = average_bc_scores.tolist()
+    mc_score_list = average_mc_scores.tolist()
+
+    # Temporary
+    mc_score_list[1] = mc_score_list[0] +  mc_score_list[1]
+    mc_score_list = mc_score_list[1:]
+    
+    bc_score = {"AI": bc_score[1].item(), "HUMAN": bc_score[0].item()}
+    mc_score = {}
+    label_map = ["OpenAI GPT", "CLAUDE", "BARD", "LLAMA 2"]
+    
+    for score, label in zip(mc_score_list, label_map):
+        mc_score[label.upper()] = score
     
     sum_prob = 1 - bc_score["HUMAN"]
     for key, value in mc_score.items():
@@ -273,9 +349,10 @@ def build_date(year, month, day):
     return f"{year}{months[month]}{day}"
 
 def len_validator(text):
-    min_chars = 350  
-    if len(text) < min_chars:
-        return f"Warning! Input length is {len(text)}. Please input a text that is greater than {min_chars} characters long. Recommended length {min_chars*2} characters." 
+    min_tokens = 128  
+    lengt = len(tokenizer.tokenize(text = text, return_tensors="pt"))
+    if  lengt < min_tokens:
+        return f"Warning! Input length is {lengt}. Please input a text that is greater than {min_tokens} tokens long. Recommended length {min_tokens*2} tokens." 
     else :
         return f"Input length is satisified." 
 
@@ -393,7 +470,7 @@ with gr.Blocks() as demo:
     d1 = today.strftime("%d/%B/%Y")
     d1 = d1.split("/")
 
-    model_list = ["GPT 3.5", "GPT 4", "CLAUDE", "BARD", "LLAMA2"]
+    model_list = ["OpenAI GPT", "CLAUDE", "BARD", "LLAMA2"]
     domain_list = ["com", "org", "net", "int", "edu", "gov", "mil"]
     gr.Markdown(
         """
@@ -557,4 +634,4 @@ with gr.Blocks() as demo:
     date_from = ""
     date_to = ""
 
-demo.launch(share=True, server_name="0.0.0.0", server_port = 80, auth=("polygraf-admin", "test@aisd"))
+demo.launch(share=True, server_name="0.0.0.0", auth=("polygraf-admin", "test@aisd"))