Update app.py

app.py

@@ -4,11 +4,12 @@ import io
 import base64
 import uuid
 import pixeltable as pxt
-from pixeltable.iterators import DocumentSplitter
 import numpy as np
+from pixeltable.iterators import DocumentSplitter
 from pixeltable.functions.huggingface import sentence_transformer
 from pixeltable.functions import openai
 from gradio.themes import Monochrome
+from huggingface_hub import HfApi, HfFolder
 
 import os
 import getpass
@@ -34,7 +35,18 @@ def create_prompt(top_k_list: list[dict], question: str) -> str:
     QUESTION:
     {question}'''
 
-def process_files(pdf_files, chunk_limit, chunk_separator):
+def validate_token(token):
+    try:
+        api = HfApi()
+        user_info = api.whoami(token=token)
+        return user_info is not None
+    except Exception:
+        return False
+
+def process_files(token, pdf_files, chunk_limit, chunk_separator):
+    if not validate_token(token):
+        return "Invalid token. Please enter a valid Hugging Face token."
+    
     # Initialize Pixeltable
     pxt.drop_dir('chatbot_demo', force=True)
     pxt.create_dir('chatbot_demo')
@@ -103,7 +115,11 @@ def process_files(pdf_files, chunk_limit, chunk_separator):
     # Extract the answer text from the API response
     t['gpt4omini'] = t.response.choices[0].message.content
 
-    return "Files processed successfully!"
+    return "Files processed successfully. You can start the discussion."
+
+def get_answer(token, msg):
+    if not validate_token(token):
+        return "Invalid token. Please enter a valid Hugging Face token."
 
 def get_answer(msg):
 
@@ -117,8 +133,8 @@ def get_answer(msg):
 
     return answer
 
-def respond(message, chat_history):
-    bot_message = get_answer(message)
+def respond(token, message, chat_history):
+    bot_message = get_answer(token, message)
     chat_history.append((message, bot_message))
     return "", chat_history
 
@@ -160,6 +176,8 @@ with gr.Blocks(theme=gr.themes.Base()) as demo:
             - The LLM formulates a response based on the provided context and the user's question.
           """)
 
+    user_token = gr.Textbox(label="Enter your Hugging Face Token", type="password")
+          
     with gr.Row():
         with gr.Column(scale=1):
             pdf_files = gr.File(label="Upload PDF Documents", file_count="multiple")
@@ -177,8 +195,8 @@ with gr.Blocks(theme=gr.themes.Base()) as demo:
             msg = gr.Textbox(label="Your Question", placeholder="Ask a question about the uploaded documents")
             submit = gr.Button("Submit")
 
-    process_button.click(process_files, inputs=[pdf_files, chunk_limit, chunk_separator], outputs=[process_output])
-    submit.click(respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
+    process_button.click(process_files, inputs=[user_token,pdf_files, chunk_limit, chunk_separator], outputs=[process_output])
+    submit.click(respond, inputs=[user_token, msg, chatbot], outputs=[msg, chatbot])
 
 if __name__ == "__main__":
     demo.launch()