Update app.py

app.py

@@ -1,15 +1,42 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer
 
 # Load the text classification model
 classifier = pipeline('text-classification', model='ardavey/bert-base-ai-generated-text')
 
+# Load the tokenizer to handle text preprocessing
+tokenizer = AutoTokenizer.from_pretrained('ardavey/bert-base-ai-generated-text')
+
+# Define a function to truncate or split the input text
+def preprocess_long_text(text, tokenizer, max_length=512):
+    # Tokenize the text
+    tokens = tokenizer.encode(text, add_special_tokens=True)
+    # Split into chunks of max_length
+    chunks = [tokens[i:i + max_length] for i in range(0, len(tokens), max_length)]
+    # Decode back to text
+    return [tokenizer.decode(chunk, skip_special_tokens=True) for chunk in chunks]
+
 # Define a function for text classification
 def classify_text(text):
-    predictions = classifier([text])
-    label = 'AI Generated Text' if predictions[0]['label'] == 'LABEL_1' else 'Human Generated Text'
-    score = predictions[0]['score']
-    return f"Prediction: {label}, Score: {score:.4f}"
+    # Preprocess the text for long input
+    chunks = preprocess_long_text(text, tokenizer)
+    
+    # Make predictions for each chunk
+    predictions = [classifier(chunk)[0] for chunk in chunks]
+    
+    # Aggregate results (you can customize this logic)
+    ai_scores = [pred['score'] for pred in predictions if pred['label'] == 'LABEL_1']
+    human_scores = [pred['score'] for pred in predictions if pred['label'] == 'LABEL_0']
+    
+    # Determine the overall prediction
+    if sum(ai_scores) > sum(human_scores):
+        label = "AI Generated Text"
+        score = sum(ai_scores) / len(ai_scores)
+    else:
+        label = "Human Generated Text"
+        score = sum(human_scores) / len(human_scores)
+    
+    return f"Prediction: {label}, Average Score: {score:.4f}"
 
 # Create a Gradio interface
 interface = gr.Interface(