import os import sqlite3 import requests import fitz # PyMuPDF import faiss import numpy as np from sentence_transformers import SentenceTransformer import gradio as gr # Configure Hugging Face API huggingface_api_url = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-70B-Instruct" huggingface_api_key = os.getenv("HUGGINGFACE_API_KEY") headers = {"Authorization": f"Bearer {huggingface_api_key}"} # Function to query Hugging Face model def query_huggingface(payload): response = requests.post(huggingface_api_url, headers=headers, json=payload) return response.json() # Function to extract text from PDF def extract_text_from_pdf(pdf_file): text = "" pdf_document = fitz.open(stream=pdf_file.read(), filetype="pdf") for page_num in range(len(pdf_document)): page = pdf_document.load_page(page_num) text += page.get_text() return text # Initialize SQLite database def init_db(): conn = sqlite3.connect('storage_warehouse.db') c = conn.cursor() c.execute(''' CREATE TABLE IF NOT EXISTS context ( id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT, content TEXT ) ''') conn.commit() conn.close() # Add context to the database def add_context(name, content): conn = sqlite3.connect('storage_warehouse.db') c = conn.cursor() c.execute('INSERT INTO context (name, content) VALUES (?, ?)', (name, content)) conn.commit() conn.close() # Retrieve context from the database def get_context(): conn = sqlite3.connect('storage_warehouse.db') c = conn.cursor() c.execute('SELECT content FROM context') context = c.fetchall() conn.close() return [c[0] for c in context] # Function to create or update the FAISS index def update_faiss_index(): contexts = get_context() if len(contexts) == 0: return None, contexts embeddings = model.encode(contexts, convert_to_tensor=True) index = faiss.IndexFlatL2(embeddings.shape[1]) index.add(embeddings.cpu().numpy()) return index, contexts # Retrieve relevant context from the FAISS index def retrieve_relevant_context(index, contexts, query, top_k=5): if index is None or len(contexts) == 0: return [] query_embedding = model.encode([query], convert_to_tensor=True).cpu().numpy() distances, indices = index.search(query_embedding, top_k) relevant_contexts = [contexts[i] for i in indices[0]] return relevant_contexts # Initialize the database and FAISS model init_db() model = SentenceTransformer('all-MiniLM-L6-v2') faiss_index, context_list = update_faiss_index() # Gradio interface def chatbot(question): relevant_contexts = retrieve_relevant_context(faiss_index, context_list, question) user_input = f"question: {question} context: {' '.join(relevant_contexts)}" response = query_huggingface({"inputs": user_input}) response_text = response.get("generated_text", "Sorry, I couldn't generate a response.") return response_text # File upload function def upload_pdf(file): context = extract_text_from_pdf(file) add_context(file.name, context) global faiss_index, context_list faiss_index, context_list = update_faiss_index() return "PDF content added to context." # Gradio interface iface = gr.Interface(fn=chatbot, inputs="text", outputs="text", title="Storage Warehouse Customer Service Chatbot") file_upload = gr.Interface(fn=upload_pdf, inputs="file", outputs="text", title="Upload PDF for Context") app = gr.TabbedInterface([iface, file_upload], ["Chatbot", "Upload PDF"]) app.launch()