semantic_search.py

# -*- coding: utf-8 -*-
"""Semantic_Search.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/17T5B7vqy-cO_AUjxJ4igRQm3WRJ5um8-
"""
import openai
# Commented out IPython magic to ensure Python compatibility.
# openai.api_key="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"
openai.api_key="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"

# !pip install openai
# !pip install tiktoken
# !pip install transformers

from fastapi import FastAPI
app = FastAPI()

# imports
import ast  # for converting embeddings saved as strings back to arrays
import openai  # for calling the OpenAI API
import pandas as pd  # for storing text and embeddings data
import tiktoken  # for counting tokens
from scipy import spatial  # for calculating vector similarities for search


# models
EMBEDDING_MODEL = "text-embedding-ada-002"
GPT_MODEL = "gpt-3.5-turbo-16k"

df_embeddings_clean = pd.read_csv('./content/embeddings_clean.csv')

"""### 0. Cleaning data with DaVinci



---


"""

@app.get("/")
def read_root():
    return {"Hello": "World"}

# def get_clean_version(context, temp=0):
#     try:
#         response = openai.ChatCompletion.create(
#             model="gpt-3.5-turbo",
#             messages=[{"role": "system", "content": f"Clean the following document from an oil drilling process report to feed a large language model considering punctuations, grammar, lowercase and uppercase, eliminate repetitive information, and add english connector words where necessary.\nText: {context}"}]
#         )
#         return response['choices'][0]['message']['content']
#     except Exception as e:
#         print(e)
#         return ""

"""### 1. Embedding the context

---


"""

def num_tokens(text: str, model: str = GPT_MODEL) -> int:
    """Return the number of tokens in a string."""
    encoding = tiktoken.encoding_for_model(model)
    return len(encoding.encode(text))

# df['tokens'] = df['clean_gpt3.5'].apply(num_tokens)

# df['tokens'].describe()

# calculate embeddings
# EMBEDDING_MODEL = "text-embedding-ada-002"  # OpenAI's best embeddings as of Apr 2023
# BATCH_SIZE = 10  # you can submit up to 2048 embedding inputs per request

# embeddings = []
# for batch_start in range(0, len(df['clean_gpt3.5']), BATCH_SIZE):
#     batch_end = batch_start + BATCH_SIZE
#     batch = df['clean_gpt3.5'][batch_start:batch_end].to_list()
#     print(f"Batch {batch_start} to {batch_end-1}")
#     response = openai.Embedding.create(model=EMBEDDING_MODEL, input=batch)
#     for i, be in enumerate(response["data"]):
#         assert i == be["index"]  # double check embeddings are in same order as input
#     batch_embeddings = [e["embedding"] for e in response["data"]]
#     embeddings.extend(batch_embeddings)

# df_embeddings_clean = pd.DataFrame({"text": df['clean_gpt3.5'], "embedding": embeddings})

# df_embeddings_clean.to_csv('./content/embeddings_clean.csv')
df_embeddings_clean['embedding'] = df_embeddings_clean['embedding'].apply(ast.literal_eval)

# df_embeddings_clean

# search function
def strings_ranked_by_relatedness(
    query: str,
    df: pd.DataFrame,
    relatedness_fn=lambda x, y: 1 - spatial.distance.cosine(x, y),
    top_n: int = 10
) -> tuple[list[str], list[float]]:
    """Returns a list of strings and relatednesses, sorted from most related to least."""
    query_embedding_response = openai.Embedding.create(
        model=EMBEDDING_MODEL,
        input=query,
    )
    query_embedding = query_embedding_response["data"][0]["embedding"]
    strings_and_relatednesses = [
        (row["text"], relatedness_fn(query_embedding, row["embedding"]))
        for i, row in df.iterrows()
    ]
    strings_and_relatednesses.sort(key=lambda x: x[1], reverse=True)
    strings, relatednesses = zip(*strings_and_relatednesses)
    return strings[:top_n], relatednesses[:top_n]

# strings, relatednesses = strings_ranked_by_relatedness('What is the name of the well on 11/23/2020?', df_embeddings_clean, top_n=5)
# for string, relatedness in zip(strings, relatednesses):
#     print(f"{relatedness=:.3f}")
#     print(string)

def num_tokens(text: str, model: str = GPT_MODEL) -> int:
    """Return the number of tokens in a string."""
    encoding = tiktoken.encoding_for_model(model)
    return len(encoding.encode(text))


def query_message(
    query: str,
    df: pd.DataFrame,
    model: str,
    token_budget: int
) -> str:
    """Return a message for GPT, with relevant source texts pulled from a dataframe."""
    strings, relatednesses = strings_ranked_by_relatedness(query, df)
    introduction = 'Use the below reports on oil drilling to answer the subsequent question. If the answer cannot be found in the reports, write "I could not find an answer. Rephrase your question."'
    question = f"\n\nQuestion: {query}"
    message = introduction
    for string in strings:
        next_article = f'\n\nContext:\n"""\n{string}\n"""'
        if (
            num_tokens(message + next_article + question, model=model)
            > token_budget
        ):
            break
        else:
            message += next_article
    return message + question

@app.get("/predict")
def ask(
    query: str
) -> str:
    df: pd.DataFrame = df_embeddings_clean
    model: str = GPT_MODEL
    token_budget: int = 15000
    print_message: bool = False
    """Answers a query using GPT and a dataframe of relevant texts and embeddings."""
    message = query_message(query, df, model=model, token_budget=token_budget)
    if print_message:
        print(message)
    messages = [
        {"role": "system", "content": "You answer questions about the oil drilling reports."},
        {"role": "user", "content": message},
    ]
    response = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        temperature=0
    )
    response_message = response["choices"][0]["message"]["content"]
    return response_message

#uvicorn app:app --port 8000 --reload

#%%