Update app.py

app.py

@@ -2,27 +2,58 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from datetime import datetime, timezone
 import os
-from huggingface_hub import login
+from huggingface_hub import login, HfApi
 from huggingface_hub.utils import RepositoryNotFoundError, HfHubHTTPError
+import requests
+
+def get_timestamp():
+    """Get current UTC datetime in specified format"""
+    return datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%SS')
+
+def format_system_info():
+    """Format system information header"""
+    return (
+        f"Current Date and Time (UTC - YYYY-MM-DD HH:MM:SS formatted): {get_timestamp()}\n"
+        f"Current User's Login: Raj-VedAI\n"
+    )
+
+def verify_model_access():
+    system_info = format_system_info()
+    try:
+        token = os.getenv("HUGGING_FACE_HUB_TOKEN")
+        if not token:
+            return False, f"{system_info}Status: No token found"
+            
+        # Method 1: Direct API check
+        api = HfApi(token=token)
+        try:
+            model_info = api.model_info("CohereLabs/c4ai-command-a-03-2025")
+            return True, f"{system_info}Status: ✅ Access granted\nModel: CohereLabs/c4ai-command-a-03-2025"
+        except Exception as e:
+            if "403" in str(e):
+                return False, f"{system_info}Status: ❌ Access denied\nPlease request access at https://huggingface.co/CohereLabs/c4ai-command-a-03-2025"
+            return False, f"{system_info}Status: ❌ Error\nDetails: {str(e)}"
+            
+    except Exception as e:
+        return False, f"{system_info}Status: ❌ Unexpected error\nDetails: {str(e)}"
 
 def initialize_model():
     try:
-        # Login to Hugging Face Hub
         token = os.getenv("HUGGING_FACE_HUB_TOKEN")
         if not token:
             return False, "No token found. Please set HUGGING_FACE_HUB_TOKEN in Space secrets.", None
             
         login(token=token)
         
-        # Initialize the model and tokenizer with token (not use_auth_token)
+        # Initialize the model and tokenizer with token
         model_id = "CohereLabs/c4ai-command-a-03-2025"
         tokenizer = AutoTokenizer.from_pretrained(
             model_id,
-            token=token  # Updated from use_auth_token
+            token=token
         )
         model = AutoModelForCausalLM.from_pretrained(
             model_id,
-            token=token  # Updated from use_auth_token
+            token=token
         )
         return True, model, tokenizer
     except RepositoryNotFoundError:
@@ -37,34 +68,27 @@ def initialize_model():
     except Exception as e:
         return False, f"Unexpected error: {str(e)}", None
 
-# Initialize model and handle potential errors
-success, result, tokenizer = initialize_model()
-if not success:
-    print(f"Error initializing model: {result}")
-else:
-    model = result
-
-def get_timestamp():
-    """Get current UTC datetime in specified format"""
-    return datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%SS')
-
-def format_system_info():
-    """Format system information header"""
-    return (
-        f"Current Date and Time (UTC - YYYY-MM-DD HH:MM:SS formatted): {get_timestamp()}\n"
-        f"Current User's Login: Raj-VedAI\n"
-    )
+def check_access_status():
+    access_granted, message = verify_model_access()
+    return message
 
 def chat(message, history):
-    if not success:
-        return [(message, f"Error: {result}")]
+    system_info = format_system_info()
+    
+    # Verify access before proceeding
+    access_granted, status_message = verify_model_access()
+    if not access_granted:
+        return [(message, f"{system_info}Error: {status_message}")]
     
     if history is None:
         history = []
     
     try:
-        # Add system information
-        system_info = format_system_info()
+        # Initialize model if not already done
+        success, result, tokenizer = initialize_model()
+        if not success:
+            return [(message, f"{system_info}Error: {result}")]
+        model = result
         
         # Format messages with the c4ai-command-a-03-2025 chat template
         messages = [{"role": "user", "content": message}]
@@ -82,23 +106,35 @@ def chat(message, history):
         gen_text = tokenizer.decode(gen_tokens[0])
         
         # Format the full response with system info
-        formatted_response = f"{system_info}\n{gen_text}"
+        formatted_response = f"{system_info}{gen_text}"
         history.append((message, formatted_response))
         return history
     except Exception as e:
-        return [(message, f"Error during chat: {str(e)}")]
+        return [(message, f"{system_info}Error during chat: {str(e)}")]
 
-# Create the Gradio interface - removed retry_on_error
-demo = gr.ChatInterface(
-    fn=chat,
-    title="Medical Decision Support AI",
-    description="""A medical decision support system that provides healthcare-related information and guidance.
-    Current UTC Time: """ + get_timestamp(),
-    examples=[
-        "What are the symptoms of hypertension?",
-        "What are common drug interactions with aspirin?",
-        "What are the warning signs of diabetes?",
-    ]
-)
+# Create the Gradio interface with both chat and status check
+with gr.Blocks(theme=gr.themes.Default()) as demo:
+    gr.Markdown(f"# Medical Decision Support AI\n{format_system_info()}")
+    
+    with gr.Row():
+        status_btn = gr.Button("Check Access Status")
+        status_output = gr.Textbox(label="Access Status", lines=6)
+    
+    chat_interface = gr.ChatInterface(
+        fn=chat,
+        description="A medical decision support system that provides healthcare-related information and guidance.",
+        examples=[
+            "What are the symptoms of hypertension?",
+            "What are common drug interactions with aspirin?",
+            "What are the warning signs of diabetes?",
+        ]
+    )
+    
+    status_btn.click(check_access_status, outputs=status_output)
+    
+    # Perform initial access check
+    access_granted, status_message = verify_model_access()
+    if not access_granted:
+        gr.Warning(status_message)
 
 demo.launch()