Update appbkup.py

appbkup.py

@@ -1,7 +1,8 @@
 import os
 from threading import Thread
 from typing import Iterator
-
+import requests
+import json
 import gradio as gr
 import spaces
 import torch
@@ -9,8 +10,8 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
 
 HF_TOKEN = "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"  ##Llama Write Token
 
-MAX_MAX_NEW_TOKENS = 2048
-DEFAULT_MAX_NEW_TOKENS = 1024
+MAX_MAX_NEW_TOKENS = 8192
+DEFAULT_MAX_NEW_TOKENS = 4096
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
 DESCRIPTION = """\
@@ -32,13 +33,14 @@ def generate(
     message: str,
     chat_history: list[tuple[str, str]],
     system_prompt: str,
-    max_new_tokens: int = 1024,
+    max_new_tokens: int = 8192,
     temperature: float = 0.6,
     top_p: float = 0.9,
     top_k: int = 50,
     repetition_penalty: float = 1.2,
 ) -> Iterator[str]:
     conversation = []
+
     if system_prompt:
         conversation.append({"role": "system", "content": system_prompt})
     for user, assistant in chat_history:
@@ -66,10 +68,45 @@ def generate(
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
 
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        yield "".join(outputs)
+    concatenated_outputs = "".join([r"{}".format(text) for text in streamer])
+
+    # Mask the output here
+    masked_output = mask_with_protecto(concatenated_outputs)
+    masked_output = format_for_html(masked_output)
+    yield masked_output
+
+
+# Ensuring entity tags are properly rendered
+def format_for_html(text):
+    text = text.replace("<", "&lt;")
+    text = text.replace(">", "&gt;")
+    return text
+
+    
+def mask_with_protecto(text_for_prompt):
+    mask_request_url = "https://trial.protecto.ai/api/vault/mask"
+    headers = {
+        "Content-Type": "application/json; charset=utf-8",
+        "Authorization": "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3N1ZXIiOiJQcm90ZWN0byIsImV4cGlyYXRpb25fZGF0ZSI6IjIwMjMtMTEtMDQiLCJwZXJtaXNzaW9ucyI6WyJyZWFkIiwid3JpdGUiXSwidXNlcl9uYW1lIjoiZGlwYXlhbkBjb2V1c2xlYXJuaW5nLmNvbSIsImRiX25hbWUiOiJwcm90ZWN0b19jb2V1c2xlYXJuaW5nX25ydG1mYmFrIiwiaGFzaGVkX3Bhc3N3b3JkIjoiMjIyMTI2ZWNiZTlkZTRmNWJlODdiY2QyYWFlZWRlM2FmNDc5MzMxZmNhOTUxMWU0MDRiNzkxNDM1MGI4MWUyYiJ9.DeIK00NuhM51lRwWdnUXuQSBA1aBn5AQ8qM3pIeM01U"
+    }
+    mask_input = {
+    "mask": [
+        {
+            "value": text_for_prompt
+        }
+    ]
+    }
+    response = requests.put(mask_request_url, headers=headers, json=mask_input)
+    if response.status_code == 200:
+        # Parse the masked result from the API response and format it for display
+        masked_result = response.json()
+        masked_result_token_value = str(masked_result["data"][0]["token_value"])
+        return_value = masked_result_token_value
+        return(return_value)
+    else:
+        # Return an error message if the API request was not successful.
+        return(str(response.status_code))
+
 
 
 chat_interface = gr.ChatInterface(