Update app.py

app.py

@@ -8,6 +8,40 @@ st.title("AI Detectability Index (ADI)")
 # Create a container for the box
 box = st.container()
 
+from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList
+
+tokenizer = AutoTokenizer.from_pretrained("StabilityAI/stablelm-tuned-alpha-7b")
+model = AutoModelForCausalLM.from_pretrained("StabilityAI/stablelm-tuned-alpha-7b")
+model.half().cuda()
+
+class StopOnTokens(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        stop_ids = [50278, 50279, 50277, 1, 0]
+        for stop_id in stop_ids:
+            if input_ids[0][-1] == stop_id:
+                return True
+        return False
+
+system_prompt = """<|SYSTEM|># StableLM Tuned (Alpha version)
+- StableLM is a helpful and harmless open-source AI language model developed by StabilityAI.
+- StableLM is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
+- StableLM is more than just an information source, StableLM is also able to write poetry, short stories, and make jokes.
+- StableLM will refuse to participate in anything that could harm a human.
+"""
+
+prompt = f"{system_prompt}<|USER|>What's your mood today?<|ASSISTANT|>"
+
+inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+tokens = model.generate(
+  **inputs,
+  max_new_tokens=64,
+  temperature=0.7,
+  do_sample=True,
+  stopping_criteria=StoppingCriteriaList([StopOnTokens()])
+)
+#print(tokenizer.decode(tokens[0], skip_special_tokens=True))
+
+
 # Create two columns inside the box
 with box:
     col1, col2 = st.columns([0.4, 0.6])
@@ -19,4 +53,4 @@ with box:
     # Add text to the right column (60% area)
     with col2:
         st.write("By conducting extensive experiments (detailed in Section 3), our study provides a thorough investigation of the de-watermarking techniques wv1 and wv2, demonstrating that the watermarked texts generated by both methods can be circumvented, albeit with a slight decrease in de-watermarking accuracy observed with wv2. These results further strengthen our contention that text watermarking is fragile and lacks reliability for real-life applications.")
-        st.write("You can add more text or components here.")
+        st.write(tokenizer.decode(tokens[0], skip_special_tokens=True))