Spaces:
Sleeping
Sleeping
File size: 5,133 Bytes
1a86be4 d77afce 1a86be4 d77afce 1a86be4 d77afce 1a86be4 d77afce 1a86be4 d77afce 1a86be4 d77afce 1a86be4 |
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 |
import os
import time
import json
import logging
import threading
import gradio as gr
import google.generativeai as genai
from googleapiclient.discovery import build
from googleapiclient.http import MediaIoBaseDownload
from google.oauth2 import service_account
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader, TextLoader, Docx2txtLoader
from langchain.chains import RetrievalQA
from langchain_google_genai import ChatGoogleGenerativeAI
from PyPDF2 import PdfReader
from gtts import gTTS
# ✅ Configure logging
logging.basicConfig(level=logging.INFO)
# ✅ Load API Keys
logging.info("🔑 Loading API keys...")
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY_1")
SERVICE_ACCOUNT_JSON = os.getenv("SERVICE_ACCOUNT_JSON")
if not GOOGLE_API_KEY or not SERVICE_ACCOUNT_JSON:
logging.error("❌ Missing API Key or Service Account JSON.")
raise ValueError("❌ Missing API Key or Service Account JSON. Please add them as environment variables.")
os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY
SERVICE_ACCOUNT_FILE = json.loads(SERVICE_ACCOUNT_JSON)
SCOPES = ["https://www.googleapis.com/auth/drive"]
FOLDER_ID = "1xqOpwgwUoiJYf9GkeuB4dayme4zJcujf"
creds = service_account.Credentials.from_service_account_info(SERVICE_ACCOUNT_FILE)
drive_service = build("drive", "v3", credentials=creds)
# ✅ Initialize variables
vector_store = None
file_id_map = {}
temp_dir = "./temp_downloads"
os.makedirs(temp_dir, exist_ok=True)
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
# ✅ Get list of files from Google Drive
def get_files_from_drive():
logging.info("📂 Fetching files from Google Drive...")
query = f"'{FOLDER_ID}' in parents and trashed = false"
results = drive_service.files().list(q=query, fields="files(id, name)").execute()
files = results.get("files", [])
global file_id_map
file_id_map = {file["name"]: file["id"] for file in files}
return list(file_id_map.keys()) if files else []
# ✅ Download file from Google Drive
def download_file(file_id, file_name):
file_path = os.path.join(temp_dir, file_name)
request = drive_service.files().get_media(fileId=file_id)
with open(file_path, "wb") as f:
downloader = MediaIoBaseDownload(f, request)
done = False
while not done:
_, done = downloader.next_chunk()
return file_path
# ✅ Process documents
def process_documents(selected_files):
global vector_store
docs = []
for file_name in selected_files:
file_path = download_file(file_id_map[file_name], file_name)
if file_name.endswith(".pdf"):
loader = PyPDFLoader(file_path)
elif file_name.endswith(".txt"):
loader = TextLoader(file_path)
elif file_name.endswith(".docx"):
loader = Docx2txtLoader(file_path)
else:
logging.warning(f"⚠️ Unsupported file type: {file_name}")
continue
docs.extend(loader.load())
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
split_docs = text_splitter.split_documents(docs)
vector_store = Chroma.from_documents(split_docs, embeddings)
return "✅ Documents processed successfully!"
# ✅ Query document
def query_document(question):
if vector_store is None:
return "❌ No documents processed.", None
retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
model = ChatGoogleGenerativeAI(model="gemini-2.0-flash", google_api_key=GOOGLE_API_KEY)
qa_chain = RetrievalQA.from_chain_type(llm=model, retriever=retriever)
response = qa_chain.invoke({"query": question})["result"]
tts = gTTS(text=response, lang="en")
temp_audio_path = os.path.join(temp_dir, "response.mp3")
tts.save(temp_audio_path)
return response, temp_audio_path
# ✅ Gradio UI
with gr.Blocks() as demo:
gr.Markdown("# 📄 AI-Powered Multi-Document Chatbot with Voice Output")
file_dropdown = gr.Dropdown(choices=get_files_from_drive(), label="📂 Select Files", multiselect=True)
refresh_button = gr.Button("🔄 Refresh Files") # 🔄 Add Refresh Button
process_button = gr.Button("🚀 Process Documents")
user_input = gr.Textbox(label="🔎 Ask a Question")
submit_button = gr.Button("💬 Get Answer")
response_output = gr.Textbox(label="📝 Response")
audio_output = gr.Audio(label="🔊 Audio Response")
# 🔄 Function to Refresh File List
def refresh_files():
return gr.update(choices=get_files_from_drive())
# ✅ Connect Refresh Button
refresh_button.click(refresh_files, outputs=file_dropdown)
# ✅ Connect Process Button
process_button.click(process_documents, inputs=file_dropdown, outputs=response_output)
# ✅ Connect Query Button
submit_button.click(query_document, inputs=user_input, outputs=[response_output, audio_output])
demo.launch()
|