Update app.py

app.py

@@ -22,23 +22,29 @@ CONFIG = synthid_mixin.DEFAULT_WATERMARKING_CONFIG
 # Function to check for AI-generated content using SynthID
 @spaces.GPU
 def check_plagiarism(text):
-    # Tokenize and process the input text
-    tokens = tokenizer.encode_plus(text, return_tensors="pt", truncation=True, padding=True)
-    tokens = tokens.to(DEVICE)
-
     # Logits processor for SynthID
     logits_processor = logits_processing.SynthIDLogitsProcessor(
         **CONFIG, top_k=40, temperature=0.5
     )
 
     
-    # Use SynthID's bayesian detector to check for AI generation likelihood
+    # Tokenize and process the input text
+    inputs = tokenizer(text, return_tensors="pt").to(DEVICE)
+    
+    # Generate output with model, capturing scores (logits)
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs['input_ids'], 
+            max_length=inputs['input_ids'].shape[1] + 50,  # Generate up to 50 additional tokens
+            output_scores=True, 
+            return_dict_in_generate=True
+        )
+    
+    # Process logits through SynthID to check for watermark presence
     try:
-        # Assuming the logits processor can be used to score watermarked content
-        logits_scores = logits_processor(tokens['input_ids'])
-        
-        # Simple threshold: assuming logits indicate watermark presence
-        is_watermarked = logits_scores.mean().item() > 0.5
+        # Pass logits (scores) to the SynthIDLogitsProcessor
+        logits = outputs.scores  # Extract logits from the generation output
+        is_watermarked = logits_processor(inputs['input_ids'], logits=logits).mean().item() > 0.5
         
         if is_watermarked:
             return "Flagged as AI-generated content (Academic Integrity Warning)."