| 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() |