Update app.py

app.py

@@ -6,59 +6,40 @@ import torch
 class ModelProcessor:
     def __init__(self, repo_id="HuggingFaceTB/cosmo-1b"):
         self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
-        # Initialize the tokenizer
         self.tokenizer = AutoTokenizer.from_pretrained(repo_id, use_fast=True)
-
-        # Initialize and configure the model
         self.model = AutoModelForCausalLM.from_pretrained(
             repo_id, torch_dtype=torch.float16, device_map={"": self.device}, trust_remote_code=True
         )
-        self.model.eval()  # Set the model to evaluation mode
-
-        # Set padding token as end-of-sequence token
+        self.model.eval()
         self.tokenizer.pad_token = self.tokenizer.eos_token
 
     @torch.inference_mode()
     def process_data_and_compute_statistics(self, prompt):
-        # Tokenize the prompt and move to the device
         tokens = self.tokenizer(
             prompt, return_tensors="pt", truncation=True, max_length=512
         ).to(self.model.device)
-
-        # Get the model outputs and logits
         outputs = self.model(tokens["input_ids"])
         logits = outputs.logits
-
-        # Shift right to align with logits' prediction position
         shifted_labels = tokens["input_ids"][..., 1:].contiguous()
         shifted_logits = logits[..., :-1, :].contiguous()
-
-        # Calculate entropy
         shifted_probs = torch.softmax(shifted_logits, dim=-1)
         shifted_log_probs = torch.log_softmax(shifted_logits, dim=-1)
         entropy = -torch.sum(shifted_probs * shifted_log_probs, dim=-1).squeeze()
-
-        # Flatten the logits and labels
         logits_flat = shifted_logits.view(-1, shifted_logits.size(-1))
         labels_flat = shifted_labels.view(-1)
-
-        # Calculate the negative log-likelihood loss
         probabilities_flat = torch.softmax(logits_flat, dim=-1)
         true_class_probabilities = probabilities_flat.gather(
             1, labels_flat.unsqueeze(1)
         ).squeeze(1)
         nll = -torch.log(
             true_class_probabilities.clamp(min=1e-9)
-        )  # Clamp to prevent log(0)
-
+        )
         ranks = (
             shifted_logits.argsort(dim=-1, descending=True)
             == shifted_labels.unsqueeze(-1)
         ).nonzero()[:, -1]
-
         if entropy.clamp(max=4).median() < 2.0:
             return 1
-
         return 1 if (ranks.clamp(max=4) * nll.clamp(max=4)).mean() < 5.2 else 0
 
 processor = ModelProcessor()
@@ -67,9 +48,9 @@ processor = ModelProcessor()
 def detect(prompt):
     prediction = processor.process_data_and_compute_statistics(prompt)
     if prediction == 1:
-        return "The text is likely **generated** by a language model."
+        return "<div class='output-text'>The text is likely <b>generated</b> by a language model.</div>"
     else:
-        return "The text is likely **not generated** by a language model."
+        return "<div class='output-text'>The text is likely <b>not generated</b> by a language model.</div>"
 
 with gr.Blocks(
     css="""
@@ -118,7 +99,7 @@ with gr.Blocks(
                 label="Prompt",
             )
             submit_btn = gr.Button("Submit", variant="primary")
-            output = gr.Markdown(elem_id="output-text")
+            output = gr.HTML()  # Changed to gr.HTML() to support custom HTML
 
     submit_btn.click(fn=detect, inputs=prompt, outputs=output)