Update app.py

app.py

@@ -7,9 +7,16 @@ from synthid_text import synthid_mixin, logits_processing
 # Configurations and model selection
 MODEL_NAME = "google/gemma-7b-it"
 DEVICE = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
+TOP_K = 40
+TOP_P = 0.99
+TEMPERATURE= 0.5
 
 # Initialize model and tokenizer
-model = transformers.AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(DEVICE)
+model = synthid_mixin.SynthIDGemmaForCausalLM.from_pretrained(
+    MODEL_NAME,
+    device_map=DEVICE,
+    torch_dtype=torch.bfloat16,
+)
 tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)
 tokenizer.pad_token = tokenizer.eos_token
 tokenizer.padding_side = "left"
@@ -24,22 +31,25 @@ CONFIG = synthid_mixin.DEFAULT_WATERMARKING_CONFIG
 def check_plagiarism(text):
     # Logits processor for SynthID
     logits_processor = logits_processing.SynthIDLogitsProcessor(
-        **CONFIG, top_k=40, temperature=0.5
+        **CONFIG, top_k=TOP_K, temperature=TEMPERATURE
     )
     
     # Tokenize and process the input text
-    inputs = tokenizer(text, return_tensors="pt").to(DEVICE)
+    inputs = tokenizer(text, return_tensors="pt", padding=True).to(DEVICE)
     inputs_len = inputs['input_ids'].shape[1]
 
     # Generate output with model, capturing scores (logits)
     with torch.no_grad():
         outputs = model.generate(
-            inputs['input_ids'], 
-            max_length=inputs_len + 50,  # Generate up to 50 additional tokens
-            output_scores=True, 
-            return_dict_in_generate=True
+            **inputs,
+            do_sample=True,
+            max_length=1024,
+            temperature=TEMPERATURE,
+            top_k=TOP_K,
+            top_p=TOP_P,
         )
-    
+
+      
     # Extract the generated tokens from the model's predictions
     generated_tokens = outputs.sequences[:, inputs_len:]