Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
@@ -8,13 +9,12 @@ For more information on `huggingface_hub` Inference API support, please check th
 
 #Update: Using a new base model
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-#client = InferenceClient("HuggingFaceH4/zephyr-7b-gemma-v0.1")
-#topic_model = BERTopic.load("MaartenGr/BERTopic_Wikipedia")
-# Train model
-#topic_model = BERTopic("english")
-#topics, probs = topic_model.fit_transform(docs)
 dataset = load_dataset("JustKiddo/KiddosVault")
 
+# Load the tokenizer and model for token display
+tokenizer = AutoTokenizer.from_pretrained("t5-small") #Google's T5 Model
+model = AutoModelForSeq2SeqLM.from_pretrained("t5-small")
+
 def respond(
     message,
     history: list[tuple[str, str]],
@@ -47,10 +47,24 @@ def respond(
         response += token
         yield response
 
+#My custom token generator
+def generate_tokens(text):
+    input = tokenizer(text, return_tensors="pt")
+    output = model.generate(**input)
+
+    input_ids = input["input_ids"].tolist()[0]
+    output_ids = output.tolist()[0]
+
+    input_tokens_str = tokenizer.convert_ids_to_tokens(input_ids)
+    output_tokens_str = tokenizer.convert_ids_to_tokens(output_ids)
+
+    return " ".join(input_tokens_str), " ".join(output_tokens_str)
+
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
 """
-demo = gr.ChatInterface(
+
+chatInterface = gr.ChatInterface(
     respond,
     additional_inputs=[
         gr.Textbox(value="You are a professional Mental Healthcare Chatbot.", label="System message"),
@@ -66,6 +80,21 @@ demo = gr.ChatInterface(
     ],
 )
 
+with gr.Blocks() as demo:
+    with gr.Row():
+        chatInterface
+    with gr.Column():
+        input_text = gr.Textbox(label="Input text")
+        input_tokens = gr.Textbox(label="Input tokens")
+        output_tokens = gr.Textbox(label="Output tokens")
+
+        def update_tokens(input_text):
+            input_tokens_str, output_tokens_str = generate_tokens(input_text)
+            return input_tokens_str, output_tokens_str
+
+        input_text.change(update_tokens, 
+                          inputs=input_text,
+                          output_tokens=[input_tokens, output_tokens])
 
 if __name__ == "__main__":
     demo.launch(debug=True)