Ask questions of your PDFs with Unstructured

!pip install "unstructured[pdf]"

import os
import json
import mysql.connector
import pandas as pd
import numpy as np

import openai
from openai.embeddings_utils import get_embedding

# OpenAI API Key
openai.api_key = os.environ["OPENAI_API_KEY"]

# SingleStore DB Connection
host=os.environ["SS_HOST"]
port=3306
username=os.environ["SS_USERNAME"]
password=os.environ["SS_PASSWORD"]
database=os.environ["SS_DATABASE"]

pdf_filename = "Employee-Handbook.pdf"

from unstructured.partition.pdf import partition_pdf
from unstructured.chunking.title import chunk_by_title

elements = partition_pdf(pdf_filename,
                         strategy="fast",
                        )

chunks = chunk_by_title(elements)

# Convert JSON output into Pandas DataFrame
data = []

for c in chunks:
    row = {}
    row['Element Type'] = type(c).__name__
    row['Filename'] = c.metadata.filename
    row['Date Modified'] = c.metadata.last_modified
    row['Filetype'] = c.metadata.filetype
    row['Page Number'] = c.metadata.page_number
    row['text'] = c.text
    data.append(row)

df = pd.DataFrame(data)

# Show the DataFrame
df.head()

# Create connection to S2 Database
cnx = mysql.connector.connect(user=username,
                              password=password,
                              host=host,
                              database=database)
cnx

# Drop the existing table
drop_cursor = cnx.cursor()
drop_query = "DROP TABLE IF EXISTS unstructured_data;"
drop_cursor.execute(drop_query)

# Create a new table
create_cursor = cnx.cursor()
create_query = ("CREATE TABLE unstructured_data ("
                "element_id INT AUTO_INCREMENT PRIMARY KEY, "
                "element_type VARCHAR(255), "
                "filename VARCHAR(255), "
                "date_modified DATETIME, "
                "filetype VARCHAR(255), "
                "page_number INT, "
                "text TEXT);")
create_cursor.execute(create_query)

cnx.commit()
drop_cursor.close()
create_cursor.close()

cursor = cnx.cursor()

# Loop through the DataFrame and insert each row into the table
for i, row in df.iterrows():
    insert_query = """INSERT INTO unstructured_data (element_type, filename, date_modified, filetype, page_number, text)
                      VALUES (%s, %s, %s, %s, %s, %s);"""
    cursor.execute(insert_query, (row['Element Type'], row['Filename'], row['Date Modified'], row['Filetype'], row['Page Number'], row['text']))

cnx.commit()
cursor.close()

cursor = cnx.cursor(buffered=True)

# Add a new column for text embedding
alter_query = "ALTER TABLE unstructured_data ADD text_embedding TEXT;"
cursor.execute(alter_query)

# Select and embed all text in table
query = "SELECT text FROM unstructured_data;"
cursor.execute(query)
rows = cursor.fetchall()

for i in rows:
    text_embedding = json.dumps(get_embedding(i[0], engine="text-embedding-ada-002"))
    update_query = ("UPDATE unstructured_data SET text_embedding = %s WHERE text = %s;")
    data = (text_embedding, i[0])
    cursor.execute(update_query, data)

cnx.commit()
cursor.close()

# User query
search_string = "What are the emergency management provisions include?"
search_embedding = get_embedding(search_string, engine="text-embedding-ada-002")
search_embedding_array = np.array(search_embedding)

cursor = cnx.cursor()

# Fetch text, type, filename, and embeddings from the unstructured_data table
query = "SELECT text, element_type, filename, text_embedding FROM unstructured_data;"
cursor.execute(query)

results = cursor.fetchall()

# Compute dot product scores
scores = []
for res in results:
    text = res[0]
    type_ = res[1]
    filename = res[2]
    embedding_str = res[3]

    if embedding_str is not None:
        embedding = json.loads(embedding_str)
        embedding_array = np.array(embedding)

        # Compute dot product for all records
        score = np.dot(search_embedding_array, embedding_array)
        scores.append((text, type_, filename, score))

# Sort by score and take the top 5
top_5 = sorted(scores, key=lambda x: x[3], reverse=True)[:5]

# Close the connection
cursor.close()
cnx.close()

# Display top-k records
top_5

if top_5:
    try:
        response = openai.ChatCompletion.create(
            model="gpt-4",
            messages=[
                {"role": "system",
                 "content": "You are a useful assistant. Use the assistant's content to answer the user's query. Summarize your answer based on the context."
                },
                {"role": "assistant", "content": str(top_5)},
                {"role": "user", "content": search_string},
            ],
            temperature=0
        )

        assistant_message = response['choices'][0]['message']['content']
        print("Assistant's Response:", assistant_message)

    except Exception as e:
        print(f"OpenAI API call failed: {e}")
else:
    print("No relevant documents found.")