Update app.py

app.py

@@ -17,6 +17,10 @@ import fitz
 from transformers import GPT2LMHeadModel, GPT2TokenizerFast
 import nltk, spacy, subprocess, torch
 import plotly.graph_objects as go
+import nltk
+
+nltk.download('punkt')
+
 from writing_analysis import (
     normalize,
     preprocess_text1,
@@ -190,6 +194,55 @@ 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)
+        encoded_segments.append(encoded_segment)
+    return encoded_segments
+
+
 def predict_bc(model, tokenizer, text):
     tokens = tokenizer(
         text, padding='max_length', truncation=True, max_length=256, return_tensors="pt"
@@ -198,8 +251,7 @@ def predict_bc(model, tokenizer, text):
     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(
@@ -208,17 +260,33 @@ def predict_mc(model, tokenizer, text):
     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):
-
+    
+    bc_scores = []
+    mc_scores = []
     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)
+    samples_len = len(split_text_allow_complete_sentences_nltk(input))
+
+    for i in samples_len:
+        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()
+    
+    bc_score = {"AI": bc_score[1].item(), "HUMAN": bc_score[0].item()}
+    mc_score = {}
+    label_map = ["GPT 3.5", "GPT 4", "CLAUDE", "BARD", "LLAMA 2"]
+    for score, label in zip(mc_score_list, label_map):
+        mc_score[label.upper()] = score.item()
     
     sum_prob = 1 - bc_score["HUMAN"]
     for key, value in mc_score.items():