ddi_chat / app.py
teddyllm's picture
Upload 4 files
4a91290 verified
raw
history blame
3.58 kB
import os
import faiss
import numpy as np
import json
import gradio as gr
from openai import OpenAI
from sentence_transformers import SentenceTransformer
# Step 1: Set up OpenAI API key
openai_api_key = os.environ.get("OPENAI_API_KEY", "")
client = OpenAI(api_key=openai_api_key)
# Step 2: Load the pre-trained FAISS index and SentenceTransformer model
index = faiss.read_index("faiss_index.bin")
model = SentenceTransformer('all-MiniLM-L6-v2')
def load_documents(docs_path):
with open(docs_path, 'r', encoding='utf-8') as file:
return json.load(file)
# Specify the path to your JSON file
docs_path = 'documents.json'
documents = load_documents(docs_path)
dimension = 1536
def get_embeddings(text):
response = client.embeddings.create(
model="text-embedding-3-small",
input = [text]
)
embedding = response.data[0].embedding
return np.array(embedding, dtype='float32')
# Step 3: Function to search FAISS index
def search_index(query, k=3):
# Convert query to an embedding
query_vector = get_embeddings(query).reshape(1, -1).astype('float32')
# Check if the index is not empty before searching
if index.ntotal == 0:
return "No documents in the index."
# Search the FAISS index for the nearest neighbors
distances, indices = index.search(query_vector, k)
# Retrieve the top matching documents
results = [documents[i] for i in indices[0] if i != -1]
if results:
return "\n\n".join(results)
else:
return "No relevant documents found."
# Step 4: Function to generate a response using OpenAI's GPT
def generate_response(context, user_input):
prompt = f"{context}\n\nUser: {user_input}\nAssistant:"
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": prompt}],
# stream=True,
)
# for chunk in stream:
# if chunk.choices[0].delta.content is not None:
# print(chunk.choices[0].delta.content, end="")
return response.choices[0].message.content
# Step 5: Gradio chatbot function
def chatbot_interface(user_input, chat_history):
# Step 5.1: Retrieve context using FAISS
context = search_index(user_input)
# Step 5.2: Generate a response using OpenAI GPT model
response = generate_response(context, user_input)
# Step 5.3: Update chat history
chat_history.append((user_input, response))
return chat_history, chat_history
def chat_gen(message, history):
history_openai_format = []
context = search_index(message)
prompt = f"{context}\n\nUser: {message}\nAssistant:"
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": prompt}],
stream=True,
)
partial_message = ""
for chunk in response:
if chunk.choices[0].delta.content is not None:
partial_message = partial_message + chunk.choices[0].delta.content
yield partial_message
initial_msg = "Hello! I am DII assistant. You can ask me anything about DDI program. I am happy to assist you."
chatbot = gr.Chatbot(value = [[None, initial_msg]])
demo = gr.ChatInterface(chat_gen, chatbot=chatbot).queue()
try:
demo.launch(debug=True, share=False, show_api=False)
demo.close()
except Exception as e:
demo.close()
print(e)
raise e