Update app.py

app.py

@@ -1,4 +1,4 @@
-# This is a project of Chakra Lab LLC. All rights reserved.
+# This is a project of Zentropi Inc. All rights reserved.
 
 import gradio as gr
 import os
@@ -10,19 +10,11 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-#base_model_name = "google/gemma-7b"
 base_model_name = "google/gemma-2-9b"
-#adapter_model_name = "samidh/cope-g7bq-2c-hs.s1.5fpc.9-sx.s1.5.9o-VL.s1.5.9-HR.s5-SH.s5-l5e5-e3-d25-r8"
-#adapter_model_name = "cope-project/cope-g2.9b-2c-hs.s1.5fpc.9-sx.s1.5.9o-vl.s1.5.9-hr.s5-sh.s5.l5e5-e3-d0-r8"
-
-#adapter_model_name = "samidh/cope-g2.9b-2c-hs.at-sx.at-vl.s5-sh.s5-hr.s5-e2"
 adapter_model_name = "samidh/cope.a.r09.2"
 
 bnb_config = BitsAndBytesConfig(
     load_in_8bit=True,
-    #bnb_4bit_quant_type="nf4",
-    #bnb_4bit_compute_dtype=torch.bfloat16,
-    #bnb_4bit_use_double_quant=True
 )
 
 model = AutoModelForCausalLM.from_pretrained(base_model_name,
@@ -120,25 +112,14 @@ def predict(content, policy):
         # Decode the predicted token
         decoded_output = tokenizer.decode([predicted_token_id])
         
-        # Get the probability of the predicted token
-        predicted_prob = probabilities[0, predicted_token_id].item()
-        
-        # Function to get probability for a specific token
-        def get_token_probability(token):
-            token_id = tokenizer.encode(token, add_special_tokens=False)[0]
-            return probabilities[0, token_id].item()
-
-        predicted_prob_0 = get_token_probability('0')
-        predicted_prob_1 = get_token_probability('1')
-        
     if decoded_output == '1':
-        return f'VIOLATING\n(P: {predicted_prob_1:.2f})'
+        return f'MATCHES (i.e., violating)'
     else:
-        return f'NON-Violating\n(P: {predicted_prob_0:.2f})'
+        return f'NO MATCHES (i.e., non-violating)'
 
 with gr.Blocks() as iface:
-    gr.Markdown("# CoPE Alpha Preview")
-    gr.Markdown("See if the given content violates your given policy.")
+    gr.Markdown("# Zentropi CoPE Demo")
+    gr.Markdown("See if the given content matches any of your given policy labels.")
     
     with gr.Row():
         content_input = gr.Textbox(label="Content", lines=2, value=DEFAULT_CONTENT)
@@ -150,15 +131,13 @@ with gr.Blocks() as iface:
     gr.Markdown("""
     ## About CoPE
     
-    CoPE (the COntent Policy Evaluation engine) is a small language model capable of accurate content policy labeling. This is a **preview** of our alpha release and is strictly for **research** purposes. This should **NOT** be used for any production use cases.
+    CoPE (the COntent Policy Evaluation engine) is a small language model capable of accurate content policy labeling. This is a **demo* of our initial release and should **NOT** be used for any production use cases.
     
     ## How to Use
                                                                                                                                                                                                                                                                        
     1. Enter your content in the "Content" box.
     2. Specify your policy in the "Policy" box.
     3. Click "Submit" to see the results.
-
-    **Note**: Inference times are **slow** (1-2 seconds) since this is built on dev infra and not yet optimized for live systems. Please be patient!
     
     ## More Info