RAGVista / app.py
zayeem00's picture
Update app.py
284d931 verified
raw
history blame
3.46 kB
import gradio as gr
import pandas as pd
import sqlite3
import openai
import pinecone
import chromadb
from PyPDF2 import PdfReader
from transformers import pipeline
import os
from google.colab import auth
from googleapiclient.discovery import build
# Initialize APIs
openai.api_key = "YOUR_OPENAI_API_KEY"
pinecone.init(api_key="YOUR_PINECONE_API_KEY", environment="us-west1-gcp")
db = chromadb.Client() # Initialize ChromaDB
# Set up Gradio interface
def initialize_models(model_choice):
if model_choice == 'OpenAI':
return {
'embedding': openai.Embedding.create,
'chat': openai.ChatCompletion.create
}
elif model_choice == 'HuggingFace':
embedding_model = pipeline('feature-extraction', model='bert-base-uncased')
chat_model = pipeline('conversational', model='facebook/blenderbot-400M-distill')
return {
'embedding': embedding_model,
'chat': chat_model
}
def fetch_pdf_from_drive(file_id):
auth.authenticate_user()
drive_service = build('drive', 'v3')
request = drive_service.files().get_media(fileId=file_id)
file = io.BytesIO(request.execute())
pdf_reader = PdfReader(file)
text = ""
for page in pdf_reader.pages:
text += page.extract_text()
return text
def query_db(query, db_type):
conn = sqlite3.connect('data.db')
cursor = conn.cursor()
if db_type == 'Pinecone':
index = pinecone.Index('your-index-name')
results = index.query(query)
return results
elif db_type == 'ChromaDB':
# Example of ChromaDB query - adapt as needed
results = db.query(query)
return results
conn.close()
def generate_response(model_choice, query, chat_history, db_type):
models = initialize_models(model_choice)
if model_choice == 'OpenAI':
response = models['chat'](model='gpt-3.5-turbo', messages=chat_history + [{'role': 'user', 'content': query}])
return response['choices'][0]['message']['content'], chat_history + [{'role': 'user', 'content': query}, {'role': 'assistant', 'content': response['choices'][0]['message']['content']}]
elif model_choice == 'HuggingFace':
response = models['chat'](query)
return response['generated_text'], chat_history + [{'role': 'user', 'content': query}, {'role': 'assistant', 'content': response['generated_text']}]
def process_input(model_choice, query, db_type, file_id):
if file_id:
pdf_text = fetch_pdf_from_drive(file_id)
query = f"{query} {pdf_text}"
response, updated_history = generate_response(model_choice, query, chat_history, db_type)
return response, updated_history
def gradio_interface():
with gr.Blocks() as demo:
with gr.Row():
model_choice = gr.Dropdown(['OpenAI', 'HuggingFace'], label='Model Choice', value='OpenAI')
db_type = gr.Dropdown(['ChromaDB', 'Pinecone'], label='Database Type', value='ChromaDB')
file_id = gr.Textbox(label='Google Drive File ID', placeholder='Enter Google Drive file ID (for PDFs)')
with gr.Row():
chat_history = gr.Chatbot()
query = gr.Textbox(label='Query')
submit_button = gr.Button('Submit')
submit_button.click(fn=process_input, inputs=[model_choice, query, db_type, file_id], outputs=[chat_history])
return demo
if __name__ == "__main__":
gradio_interface().launch()