Update app.py

app.py

@@ -1,9 +1,36 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from huggingface_hub import login
+from huggingface_hub import login, HfApi
+import os
+from dotenv import load_dotenv
+from flask import flash, redirect, url_for, session
+
+# Load environment variables
+load_dotenv()
 
 # Login to Hugging Face (you'll need to set HUGGINGFACE_TOKEN in your Secrets)
-login()
+def login():
+    # Get the Hugging Face token from environment variable
+    hf_token = os.getenv('HUGGING_FACE_TOKEN')
+    
+    if not hf_token:
+        flash('Hugging Face token not found. Please set the HUGGING_FACE_TOKEN environment variable.', 'error')
+        return redirect(url_for('index'))
+
+    # Use the token for authentication
+    try:
+        api = HfApi()
+        user = api.whoami(token=hf_token)
+        if user:
+            session['logged_in'] = True
+            flash('Successfully authenticated with Hugging Face', 'success')
+        else:
+            session['logged_in'] = False
+            flash('Failed to authenticate with Hugging Face', 'error')
+    except Exception as e:
+        session['logged_in'] = False
+        flash(f'Error during Hugging Face authentication: {str(e)}', 'error')
+        return redirect(url_for('index'))
 
 # Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-70B")