Spaces:
Running
Running
File size: 4,814 Bytes
6357e8e 0ccdb83 6357e8e 0ccdb83 6357e8e d9bb5f4 6357e8e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 |
import os
import gradio as gr
import torch
from langchain.chains import ConversationalRetrievalChain
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.llms import HuggingFacePipeline
from transformers import pipeline
# Set Hugging Face Cache Directory
os.environ["HF_HOME"] = "/tmp/huggingface_cache"
# Check for GPU availability
DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
# Global variables
conversation_retrieval_chain = None
chat_history = []
llm_pipeline = None
embeddings = None
persist_directory = "/tmp/chroma_db" # Storage for vector DB
def init_llm():
"""Initialize LLM and Embeddings"""
global llm_pipeline, embeddings
hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
if not hf_token:
raise ValueError("HUGGINGFACEHUB_API_TOKEN is not set in environment variables.")
model_id = "tiiuae/falcon-7b-instruct"
hf_pipeline = pipeline("text-generation", model=model_id, device=DEVICE)
llm_pipeline = HuggingFacePipeline(pipeline=hf_pipeline)
embeddings = HuggingFaceEmbeddings(
model_name="sentence-transformers/all-MiniLM-L6-v2",
model_kwargs={"device": DEVICE}
)
import time
def process_document(file):
global conversation_retrieval_chain
if not llm_pipeline or not embeddings:
init_llm()
start_time = time.time()
print(f"📂 Uploading PDF: {file.name}")
try:
# ✅ Ensure file is saved correctly
file_path = os.path.join("/tmp/uploads", file.name)
with open(file_path, "wb") as f:
f.write(file.read())
print(f"✅ PDF saved at {file_path} in {time.time() - start_time:.2f}s")
# ✅ Load PDF
start_time = time.time()
loader = PyPDFLoader(file_path)
documents = loader.load()
print(f"✅ PDF loaded in {time.time() - start_time:.2f}s")
# ✅ Split text
start_time = time.time()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=50)
texts = text_splitter.split_documents(documents)
print(f"✅ Text split in {time.time() - start_time:.2f}s")
# ✅ Create ChromaDB
start_time = time.time()
db = Chroma.from_documents(texts, embedding=embeddings, persist_directory="/tmp/chroma_db")
print(f"✅ ChromaDB created in {time.time() - start_time:.2f}s")
# ✅ Create retrieval chain
conversation_retrieval_chain = ConversationalRetrievalChain.from_llm(
llm=llm_pipeline, retriever=db.as_retriever()
)
print("✅ Document processing complete!")
return "📄 PDF uploaded and processed successfully! You can now ask questions."
except Exception as e:
print(f"❌ Error processing PDF: {str(e)}")
return f"Error: {str(e)}"
def process_prompt(prompt, chat_history_display):
"""Generate a response using the retrieval chain"""
global conversation_retrieval_chain, chat_history
if not conversation_retrieval_chain:
return chat_history_display + [("❌ No document uploaded.", "Please upload a PDF first.")]
output = conversation_retrieval_chain({"question": prompt, "chat_history": chat_history})
answer = output["answer"]
chat_history.append((prompt, answer))
return chat_history
# Define Gradio UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown("<h1 style='text-align: center;'>Personal Data Assistant</h1>")
with gr.Row():
dark_mode = gr.Checkbox(label="🌙 Toggle light/dark mode")
with gr.Column(): # ✅ Replace `gr.Box()` with `gr.Column()`
gr.Markdown("Hello there! I'm your friendly data assistant, ready to answer any questions regarding your data. Could you please upload a PDF file for me to analyze?")
file_input = gr.File(label="Upload File")
upload_button = gr.Button("📂 Upload File")
status_output = gr.Textbox(label="Status", interactive=False)
chat_history_display = gr.Chatbot(label="Chat History")
with gr.Row():
user_input = gr.Textbox(placeholder="Type your message here...", scale=4)
submit_button = gr.Button("📩", scale=1)
clear_button = gr.Button("🔄", scale=1)
# Button Click Actions
upload_button.click(process_document, inputs=file_input, outputs=status_output)
submit_button.click(process_prompt, inputs=[user_input, chat_history_display], outputs=chat_history_display)
clear_button.click(lambda: [], outputs=chat_history_display)
# Launch Gradio App
if __name__ == "__main__":
demo.launch(share=True)
|