Update app.py

app.py

@@ -3,7 +3,7 @@ from huggingface_hub import InferenceClient
 import time
 from collections import defaultdict
 
-# Model definitions
+# Model definitions with real models
 model_value = {
     "Lake 1": "google/mt5-base",
     "Lake 1 Flash": "google/gemma-2-2b-it",
@@ -17,11 +17,11 @@ access_codes = {
     "8tj82-2UvU-8Lft-Dupb": "plus"
 }
 
-# Model access levels
+# Model access levels with real models
 model_access_levels = {
     "google/mt5-base": "everyone",
-    "google/gemma-2-2b-it": "plus",  # Corrected to a string
-    "google/mt5-large": "plus",       # Corrected to a string
+    "google/gemma-2-2b-it": "plus",
+    "google/mt5-large": "plus",
     "google-bert/bert-base-multilingual-cased": "pro"
 }
 
@@ -38,69 +38,71 @@ def recommend_model(current_model, access_level):
         return "Consider upgrading to Lake 1 Pro for advanced features."
     return None
 
-def can_use_flash_model(user_id):
+def can_use_flash_model():
     current_time = time.time()
-    usage_tracker[user_id] = [t for t in usage_tracker[user_id] if current_time - t < 5 * 3600]
+    for user in usage_tracker.keys():
+        usage_tracker[user] = [t for t in usage_tracker[user] 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
+    for user in usage_tracker.keys():
+        if len(usage_tracker[user]) < 20:
+            usage_tracker[user].append(current_time)
+            return True
+    return False
+
+def respond(message, history, model):
+    messages = []
 
-def respond(message, history, model, access_level):
-    # Build the message history as a single string
-    full_message = ""
     for val in history:
         if val[0]:
-            full_message += f":User  {val[0]}\n"  # User message
+            messages.append({"role": "user", "content": val[0]})    
         if val[1]:
-            full_message += f"Assistant: {val[1]}\n"  # Assistant response
+            messages.append({"role": "assistant", "content": val[1]})
 
-    # Add the current user message
-    full_message += f":User  {message}\n"
+    messages.append({"role": "user", "content": message})
 
     client = InferenceClient(model)
     response = ""
-    try:
-        # Call the chat completion method with the correct format
-        response_message = client.chat_completion(
-            model=model,
-            inputs=full_message,  # Use the concatenated string
-            max_tokens=512,
-            temperature=0.7,
-            top_p=0.95,
-            stream=False  # Disable streaming
-        )
-        response = response_message['choices'][0]['text']  # Access the content directly
-    except Exception as e:
-        return f"Error during model inference: {str(e)}", history
+    for message in client.chat_completion(
+        messages,
+        max_tokens=512,
+        stream=True,
+        temperature=0.7,
+        top_p=0.95,
+    ):
+        token = message.choices[0].delta.content
+        response += token
 
     history.append((message, response))
     return response, history
 
-def main(message, history, model_name, access_code, user_id):
+def main(message, history, model_name, access_code):
     model = model_value[model_name]
     access_level = check_access_code(access_code)
     
+    # Check if the model is Lake 1 Flash
     if model == model_value["Lake 1 Flash"]:
-        if not can_use_flash_model(user_id):
+        if not can_use_flash_model():
             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.get(model, None)
-        if access_level == required_access_level:
-            return respond(message, history, model, access_level)
-        else:
-            recommendation = recommend_model(model_name, access_level)
-            if recommendation:
-                return f"You do not have access to the {model_name}. {recommendation} Please enter a valid access code for this model.", history
-            else:
-                return f"You do not have access to the {model_name}. Please enter a valid access code for this model.", history
+    # Get the required access level for the selected model
+    required_access_level = model_access_levels.get(model, None)
+    
+    # Allow access to Lake 1 for everyone
+    if required_access_level == "everyone":
+        return respond(message, history, model)
+
+    # Check access for other models
+    if access_level == "pro" or (access_level == "plus" and required_access_level in ["plus", "everyone"]):
+        return respond(message, history, model)
+    
+    # If the user does not have the required access, provide a recommendation
+    recommendation = recommend_model(model_name, access_level)
+    if recommendation:
+        return f"You do not have access to the {model_name}. {recommendation} Please enter a valid access code for this model.", history
     else:
-        return "Invalid access code. Please enter a valid code to use this service .", history
+        return f"You do not have access to the {model_name}. Please enter a valid access code for this model.", history
+
+    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")
@@ -109,7 +111,6 @@ with gr.Blocks() as demo:
     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():
@@ -128,7 +129,7 @@ with gr.Blocks() as demo:
 
     submit_button.click(
         fn=main,
-        inputs=[message_input, history, model_dropdown, access_code_input, user_id_input],
+        inputs=[message_input, history, model_dropdown, access_code_input],
         outputs=[response_output, history]
     )