Update app.py

app.py

@@ -1,10 +1,33 @@
 import gradio as gr
+from huggingface_hub import login, logout, whoami
 
-# Function to handle model inference with config
-def generate_text(prompt, temperature, max_tokens):
-    # Simulate model inference with config (replace with actual model call)
-    response = f"Response to '{prompt}' with temperature={temperature} and max_tokens={max_tokens}"
-    return response
+# Function to handle login
+def handle_login(token):
+    try:
+        # Attempt to log in with the provided token
+        login(token=token, add_to_git_credential=False)
+        user_info = whoami()
+        return f"Logged in as: {user_info['name']}"
+    except Exception as e:
+        # Handle login failure
+        logout()  # Ensure the user is logged out if login fails
+        return f"Login failed: {str(e)}"
+
+# Function to check if the user is logged in
+def is_logged_in():
+    try:
+        # Check if the user is authenticated
+        whoami()
+        return True
+    except:
+        return False
+
+# Function to restrict access to the app
+def restricted_functionality(prompt):
+    if not is_logged_in():
+        return "Please log in to use this feature."
+    # Simulate model response (replace with actual model inference)
+    return f"Model response to: {prompt}"
 
 # Gradio interface
 with gr.Blocks() as demo:
@@ -15,13 +38,8 @@ with gr.Blocks() as demo:
         login_button = gr.Button("Sign in")
         login_status = gr.Markdown("")
 
-        # Model configuration
-        gr.Markdown("### Model Configuration")
-        temperature = gr.Slider(0.1, 1.0, value=0.7, label="Temperature")
-        max_tokens = gr.Slider(10, 500, value=100, label="Max Tokens")
-
-    # Input and output components
-    with gr.Column():
+    # Main app functionality
+    with gr.Column(visible=False) as main_interface:  # Hide until logged in
         prompt = gr.Textbox(label="Your Prompt")
         output = gr.Textbox(label="Model Response")
         generate_button = gr.Button("Generate")
@@ -29,18 +47,18 @@ with gr.Blocks() as demo:
     # Load the model
     model_interface = gr.load("models/deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct", provider="nebius")
 
-    # Handle login (example logic)
-    def handle_login(token):
-        if token:  # Replace with actual authentication logic
-            return "Logged in successfully!"
+    # Handle login
+    def update_interface(token):
+        login_result = handle_login(token)
+        if "Logged in as:" in login_result:
+            return {main_interface: gr.update(visible=True), login_status: login_result}
         else:
-            return "Please enter a valid token."
+            return {main_interface: gr.update(visible=False), login_status: login_result}
 
-    # Handle text generation
-    generate_button.click(generate_text, [prompt, temperature, max_tokens], output)
+    login_button.click(update_interface, inputs=token_input, outputs=[main_interface, login_status])
 
-    # Handle login
-    login_button.click(handle_login, inputs=token_input, outputs=login_status)
+    # Handle text generation (restricted to logged-in users)
+    generate_button.click(restricted_functionality, inputs=prompt, outputs=output)
 
 # Launch the app
 demo.launch()