Update app.py

app.py

@@ -36,38 +36,39 @@ def test_stream(text):
 
     try:
 
-        print("39")
+        token_ids = tokenizer(text, return_tensors="pt").input_ids[0]
 
-        inputs = tokenizer(text, return_tensors="pt")
-        print("42")
-        token_ids = inputs.input_ids[0]
-        print(f"token_ids: {token_ids}")
-    
         stream_state = None
-        labels = []
-        tokens = tokenizer.convert_ids_to_tokens(token_ids)
-
-        print(f"50: {tokens}")
+        results = []
     
-        for i in range(len(token_ids)):
-
-            print("54")
+        first_token = token_ids[:1]
     
-            partial = token_ids[i]
-
-            print(f"58: {partial}")
+        result, stream_state = model.stream_moderate_from_ids(
+            first_token,
+            role="user",
+            stream_state=None
+        )
+    
+        results.append(
+            (tokenizer.decode(first_token), result["risk_level"][-1])
+        )
+    
+        for i in range(1, len(token_ids)):
+            current_token = token_ids[i]
     
             result, stream_state = model.stream_moderate_from_ids(
-                partial,
-                role="assistant",
+                current_token,
+                role="user",
                 stream_state=stream_state
             )
-
-            print("64")
     
-            labels.append(result["risk_level"][-1])
+            results.append(
+                (tokenizer.decode([current_token]), result["risk_level"][-1])
+            )
+    
+        model.close_stream(stream_state)
     
-        return "\n".join([f"{t} -> {l}" for t, l in zip(tokens, labels)])
+        return results
 
     except Exception as e:
         return str(e)