Update app.py

app.py

@@ -1,36 +1,52 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+import time
+from collections import defaultdict
 
+# Model definitions
 model_value = {
-    "Lake 1": "WhiteRabbitNeo/Trinity-13B",
-    "Lake 1 Plus": "WhiteRabbitNeo/Trinity-33B-v1.0",
-    "Lake 1 Pro": "WhiteRabbitNeo/WhiteRabbitNeo-33B-v1.5"
+    "Lake 1": "google/mt5-base",
+    "Lake 1 Flash": "google/gemma-2-2b-it",
+    "Lake 1 Plus": "google/mt5-large",
+    "Lake 1 Pro": "google-bert/bert-base-multilingual-cased"
 }
 
-client = InferenceClient(list(model_value.values()))
-
+# Access codes for different models
 access_codes = {
     "wF99-zXDg-WRiN-qVp8": "pro",
     "8tj82-2UvU-8Lft-Dupb": "plus"
 }
 
+# Model access levels
 model_access_levels = {
     "WhiteRabbitNeo/Trinity-13B": "everyone",
     "WhiteRabbitNeo/Trinity-33B-v1.0": "plus",
     "WhiteRabbitNeo/WhiteRabbitNeo-33B-v1.5": "pro"
 }
 
+# Usage tracking for "Lake 1 Flash"
+usage_tracker = defaultdict(list)
+
 def check_access_code(code):
     return access_codes.get(code, None)
 
 def recommend_model(current_model, access_level):
-    # Suggest a newer model based on the current model and access level
     if current_model == "Lake 1" and access_level == "everyone":
         return "Consider upgrading to Lake 1 Plus for more features."
     elif current_model == "Lake 1 Plus" and access_level == "plus":
         return "Consider upgrading to Lake 1 Pro for advanced features."
     return None
 
+def can_use_flash_model(user_id):
+    current_time = time.time()
+    usage_tracker[user_id] = [t for t in usage_tracker[user_id] if current_time - t < 5 * 3600]
+    
+    if len(usage_tracker[user_id]) < 20:
+        usage_tracker[user_id].append(current_time)
+        return True
+    else:
+        return False
+
 def respond(message, history, model, access_level):
     messages = []
 
@@ -42,8 +58,7 @@ def respond(message, history, model, access_level):
 
     messages.append({"role": "user", "content": message})
 
-    client.model = model
-
+    client = InferenceClient(model)
     response = ""
     for message in client.chat_completion(
         messages,
@@ -58,14 +73,18 @@ def respond(message, history, model, access_level):
     history.append((message, response))
     return response, history
 
-def main(message, history, model_name, access_code):
+def main(message, history, model_name, access_code, user_id):
     model = model_value[model_name]
     access_level = check_access_code(access_code)
     
+    if model == model_value["Lake 1 Flash"]:
+        if not can_use_flash_model(user_id):
+            return "Usage limit reached for Lake 1 Flash. Please try again later.", history
+    
     if model == model_value["Lake 1"]:
         return respond(message, history, model, "everyone")
     elif access_level:
-        required_access_level = model_access_levels[model]
+        required_access_level = model_access_levels.get(model, None)
         if access_level == required_access_level:
             return respond(message, history, model, access_level)
         else:
@@ -78,9 +97,13 @@ def main(message, history, model_name, access_code):
         return "Invalid access code. Please enter a valid code to use this service.", history
 
 with gr.Blocks() as demo:
+    gr.Markdown("## Welcome to the Model Interaction App")
+    gr.LoginButton()  # Add Hugging Face login button
+
     with gr.Row():
         with gr.Column():
             access_code_input = gr.Textbox(label="Access Code", placeholder="Enter your special access code (if needed)")
+            user_id_input = gr.Textbox(label="User  ID", placeholder="Enter your user ID")
 
     with gr.Row():
         with gr.Column():
@@ -89,17 +112,18 @@ with gr.Blocks() as demo:
                 choices=list(model_value.keys()),
                 value="Lake 1"  # Changed default value to match model_value keys
             )
-            response_output = gr.Textbox(label="Response", interactive=False, lines=10)
-            message_input = gr.Textbox(label="Message", placeholder="Type your message here...")
-            submit_button = gr.Button("Send")
+            response_output = gr.Textbox(label="Response", interactive=False, placeholder="Response will appear here")
+
+    with gr.Row():
+        message_input = gr.Textbox(label="Message", placeholder="Type your message here")
+        submit_button = gr.Button("Submit")
 
-    history_state = gr.State()
+    history = gr.State([])
 
     submit_button.click(
-        main,
-        inputs=[message_input, history_state, model_dropdown, access_code_input],
-        outputs=[response_output, history_state]
+        fn=main,
+        inputs=[message_input, history, model_dropdown, access_code_input, user_id_input],
+        outputs=[response_output, history]
     )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()