|
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 |
|
|
|
|
|
openai.api_key = "YOUR_OPENAI_API_KEY" |
|
pinecone.init(api_key="YOUR_PINECONE_API_KEY", environment="us-west1-gcp") |
|
db = chromadb.Client() |
|
|
|
|
|
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': |
|
|
|
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() |