Update app.py

app.py

@@ -1,83 +1,87 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from huggingface_hub import InferenceClient, login
 import tensorflow as tf
-from huggingface_hub import login, create_repo, upload_file
-import os
 from transformers import AutoTokenizer, TFAutoModelForCausalLM
+import os
+
+# Set up mixed precision and distribution strategy
 policy = tf.keras.mixed_precision.Policy('mixed_bfloat16')
 tf.keras.mixed_precision.set_global_policy(policy)
 strategy = tf.distribute.MultiWorkerMirroredStrategy()
 
+# Log into Hugging Face
 login(os.environ.get("hf_token"))
+
+# Load tokenizer and model
 name = "WICKED4950/GPT2mini-InstEsther0.5eV3.2"
 tokenizer = AutoTokenizer.from_pretrained(name)
 tokenizer.pad_token = tokenizer.eos_token
+
 with strategy.scope():
     model = TFAutoModelForCausalLM.from_pretrained(name)
 
+# Raw Prediction Function
 def raw_pred(input, model, tokenizer, max_length=1024, temperature=0.2):
     input_ids = tokenizer.encode(input, return_tensors='tf')
 
     # Initialize variables
     generated_ids = input_ids
-    stop_token_id = tokenizer.encode("<|SOH|>", add_special_tokens=False)[0]  # ID for <|SOH|>
+    stop_token_id = tokenizer.encode("<|SOH|>", add_special_tokens=False)[0]
     all_generated_tokens = []  # To store generated token IDs
     tokens_yielded = []  # To store tokens as they are yielded
 
     with strategy.scope():
         for _ in range(max_length // 1):  # Generate in chunks of 3 tokens
-            # Generate three tokens at a time
+            # Generate tokens
             outputs = model.generate(
                 generated_ids,
-                max_length=generated_ids.shape[1] + 1,  # Increment max length by 3
+                max_length=generated_ids.shape[1] + 1,
                 temperature=temperature,
                 pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=stop_token_id,  # Stop generation at <|SOH|>
+                eos_token_id=stop_token_id,
                 do_sample=True,
                 num_return_sequences=1
             )
 
-            # Get the newly generated tokens (last 3 tokens)
+            # Get the newly generated tokens
             new_tokens = outputs[0, -1:]
             generated_ids = outputs  # Update the generated_ids with the new tokens
 
-            # Store the generated tokens as numbers (IDs)
+            # Store and yield the generated tokens
             all_generated_tokens.extend(new_tokens.numpy().tolist())
-
-            # Decode and yield the tokens as they are generated (as numbers)
             tokens_text = tokenizer.decode(new_tokens, skip_special_tokens=False)
             tokens_yielded.append(tokens_text)
             yield tokens_text
 
-            # Stop if the generated tokens include <|SOH|>
+            # Stop if stop token is encountered
             if stop_token_id in new_tokens.numpy():
                 final_text = tokenizer.decode(all_generated_tokens, skip_special_tokens=False)
-                yield ("<|Clean|>" + final_text)
+                yield "<|Clean|>" + final_text
                 break
 
+# Response Handler Function
 def respond(message, history):
-    # Prepare input for the model
     give_mod = ""
-    history = history[-3:]
+    history = history[-3:]  # Limit history to last 3 exchanges
     for chunk in history:
-        give_mod = give_mod + "<|SOH|>" + chunk[0] + "<|SOB|>" + chunk[1]
-    give_mod = give_mod + "<|SOH|>" + message.capitalize() + "<|SOB|>"
-    print(give_mod)
+        give_mod += f"<|SOH|>{chunk[0]}<|SOB|>{chunk[1]}"
+    give_mod += f"<|SOH|>{message.capitalize()}<|SOB|>"
+
     response = ""
     for token in raw_pred(give_mod, model, tokenizer):
         if "<|Clean|>" in token:
             response = token
         else:
             response += token
-        yield response.replace("<|SOH|>","").replace("<|Clean|>","")
-    print(response)
+        yield response.replace("<|SOH|>", "").replace("<|Clean|>", "")
+
 # Gradio Chat Interface Setup
 demo = gr.ChatInterface(
-    fn=respond,  # Response handler function
-    title="Chat with Esther",  # Add a title
-    description="A friendly chatbot ready to help and chat with you! 😊",  # Brief description
-    theme="compact",  # Options: "compact", "default", "dark"
+    fn=respond,
+    title="Chat with Esther",  # Title of the app
+    description="A friendly chatbot ready to help and chat with you! 😊",  # Description of the app
+    theme="compact",  # Choose the theme
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(debug=True)