updated app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import gradio as gr
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import Chroma
@@ -9,6 +10,7 @@ import torch
 import re
 import transformers
 import spaces
+import requests
 
 # Initialize embeddings and ChromaDB
 model_name = "sentence-transformers/all-mpnet-base-v2"
@@ -24,23 +26,13 @@ books_db_client = books_db.as_retriever()
 
 # Initialize the model and tokenizer
 model_name = "stabilityai/stablelm-zephyr-3b"
-
-# bnb_config = transformers.BitsAndBytesConfig(
-#     load_in_4bit=True,
-#     bnb_4bit_quant_type='nf4',
-#     bnb_4bit_use_double_quant=True,
-#     bnb_4bit_compute_dtype=torch.bfloat16
-# )
-
 model_config = transformers.AutoConfig.from_pretrained(model_name, max_new_tokens=1024)
 model = transformers.AutoModelForCausalLM.from_pretrained(
     model_name,
     trust_remote_code=True,
     config=model_config,
-    # quantization_config=bnb_config,
     device_map=device,
 )
-
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 query_pipeline = transformers.pipeline(
@@ -70,8 +62,6 @@ books_db_client_retriever = RetrievalQA.from_chain_type(
 @spaces.GPU(duration=60)
 def test_rag(query):
     books_retriever = books_db_client_retriever.run(query)
-    
-    # Extract the relevant answer using regex
     corrected_text_match = re.search(r"Helpful Answer:(.*)", books_retriever, re.DOTALL)
     
     if corrected_text_match:
@@ -81,6 +71,14 @@ def test_rag(query):
     
     return corrected_text_books
 
+# OAuth Login Functionality
+def oauth_login():
+    client_id = os.getenv("OAUTH_CLIENT_ID")
+    redirect_uri = f"https://{os.getenv('SPACE_HOST')}/login/callback"
+    state = "random_string"  # Ideally generate a secure random string
+    login_url = f"https://huggingface.co/oauth/authorize?redirect_uri={redirect_uri}&scope=openid%20profile&client_id={client_id}&state={state}"
+    return login_url
+
 # Define the Gradio interface
 def chat(query, history=None):
     if history is None:
@@ -101,11 +99,12 @@ with gr.Blocks() as interface:
     
     input_box = gr.Textbox(label="Enter your question", placeholder="Type your question here...")
     submit_btn = gr.Button("Submit")
-    # clear_btn = gr.Button("Clear")
     chat_history = gr.Chatbot(label="Chat History")
-    
-    submit_btn.click(chat, inputs=[input_box, chat_history], outputs=[chat_history, input_box])
-    # clear_btn.click(clear_input, outputs=input_box)
 
-interface.launch()
+    # Sign-In Button
+    login_btn = gr.Button("Sign In with HF")
+    login_btn.click(lambda: oauth_login(), outputs=None)  # Redirect user for OAuth login
+
+    submit_btn.click(chat, inputs=[input_box, chat_history], outputs=[chat_history, input_box])
 
+interface.launch()