Update app.py

app.py

@@ -1,90 +1,46 @@
-
-
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from datetime import datetime, timezone
 import os
 from huggingface_hub import login, HfApi
 from huggingface_hub.utils import RepositoryNotFoundError, HfHubHTTPError
+import time
 import requests
+from tenacity import retry, stop_after_attempt, wait_exponential
 
-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)}"
-
+# Add retry decorator for connection attempts
+@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
 def initialize_model():
     try:
         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)
+        # Force re-login to refresh connection
+        login(token=token, add_to_git_credential=False)
         
-        # Initialize the model and tokenizer with token
+        # Initialize with device mapping and low memory settings
         model_id = "CohereLabs/c4ai-command-a-03-2025"
         tokenizer = AutoTokenizer.from_pretrained(
             model_id,
-            token=token
+            token=token,
+            use_fast=True
         )
         model = AutoModelForCausalLM.from_pretrained(
             model_id,
-            token=token
+            token=token,
+            device_map="auto",
+            low_cpu_mem_usage=True,
+            torch_dtype="auto"
         )
         return True, model, tokenizer
-    except RepositoryNotFoundError:
-        return False, "Model repository not found. Please check the model ID.", None
-    except HfHubHTTPError as e:
-        if e.response.status_code == 401:
-            return False, "Authentication failed. Please check your token permissions.", None
-        elif e.response.status_code == 403:
-            return False, "Access denied. Please request access at https://huggingface.co/CohereLabs/c4ai-command-a-03-2025", None
-        else:
-            return False, f"An error occurred: {str(e)}", None
     except Exception as e:
-        return False, f"Unexpected error: {str(e)}", None
-
-def check_access_status():
-    access_granted, message = verify_model_access()
-    return message
+        return False, f"Error during initialization: {str(e)}", None
 
+@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
 def chat(message, history):
     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:
         # Initialize model if not already done
         success, result, tokenizer = initialize_model()
@@ -92,35 +48,59 @@ def chat(message, history):
             return [(message, f"{system_info}Error: {result}")]
         model = result
         
-        # Format messages with the c4ai-command-a-03-2025 chat template
+        if history is None:
+            history = []
+        
+        # Format messages with the chat template
         messages = [{"role": "user", "content": message}]
         input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True)
         
-        # Generate response
+        # Generate response with safety settings
         gen_tokens = model.generate(
             input_ids,
             max_new_tokens=100,
             do_sample=True,
             temperature=0.3,
+            pad_token_id=tokenizer.eos_token_id,
+            attention_mask=input_ids.new_ones(input_ids.shape)
         )
         
         # Decode response
-        gen_text = tokenizer.decode(gen_tokens[0])
+        gen_text = tokenizer.decode(gen_tokens[0], skip_special_tokens=True)
         
         # Format the full response with system info
         formatted_response = f"{system_info}{gen_text}"
         history.append((message, formatted_response))
         return history
     except Exception as e:
-        return [(message, f"{system_info}Error during chat: {str(e)}")]
+        error_msg = f"{system_info}Error during chat: {str(e)}\nAttempting reconnection..."
+        if history is None:
+            history = []
+        history.append((message, error_msg))
+        return history
+
+def check_connection():
+    timestamp = get_timestamp()
+    try:
+        token = os.getenv("HUGGING_FACE_HUB_TOKEN")
+        api = HfApi(token=token)
+        model_info = api.model_info("CohereLabs/c4ai-command-a-03-2025")
+        return f"""
+{format_system_info()}
+Connection Status: ✅ Connected
+Model: {model_info.modelId}
+Last Modified: {model_info.lastModified}
+"""
+    except Exception as e:
+        return f"{format_system_info()}Connection Status: ❌ Error\nDetails: {str(e)}"
 
-# Create the Gradio interface with both chat and status check
+# Create the Gradio interface with connection monitoring
 with gr.Blocks(theme=gr.themes.Default()) as demo:
-    gr.Markdown(f"# Medical Decision Support AI\n{format_system_info()}")
+    gr.Markdown(f"# Medical Decision Support AI")
     
     with gr.Row():
-        status_btn = gr.Button("Check Access Status")
-        status_output = gr.Textbox(label="Access Status", lines=6)
+        connection_btn = gr.Button("Check Connection Status")
+        connection_status = gr.Textbox(label="Connection Status", lines=6)
     
     chat_interface = gr.ChatInterface(
         fn=chat,
@@ -132,11 +112,20 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
         ]
     )
     
-    status_btn.click(check_access_status, outputs=status_output)
+    connection_btn.click(check_connection, outputs=connection_status)
     
-    # Perform initial access check
-    access_granted, status_message = verify_model_access()
-    if not access_granted:
-        gr.Warning(status_message)
+    # Check connection on startup
+    connection_status.value = check_connection()
+
+# Add requirements to requirements.txt
+requirements = """
+gradio>=3.50.2
+transformers
+torch
+accelerate
+huggingface_hub
+requests
+tenacity
+"""
 
 demo.launch()