deepnote update

app.py

@@ -4,6 +4,7 @@ import faq as faq
 import util as util
 import uvicorn
 import gradio as gr
+from typing import List, Optional
 
 app = FastAPI()
 
@@ -15,6 +16,15 @@ class AskRequest(BaseModel):
     k: int
 
 
+class AskRequestEx(BaseModel):
+    question: str
+    sheet_url: str
+    page_content_column: str
+    k: int
+    id_column: str
+    synonyms: Optional[List[List[str]]] = None
+
+
 @app.post("/api/v1/ask")
 async def ask_api(request: AskRequest):
     return ask(
@@ -23,12 +33,14 @@ async def ask_api(request: AskRequest):
 
 
 @app.post("/api/v2/ask")
-async def ask_api(request: AskRequest):
+async def ask_api(request: AskRequestEx):
     util.SPLIT_PAGE_BREAKS = True
+    if request.synonyms is not None:
+        util.SYNONYMS = request.synonyms
     vectordb = faq.load_vectordb(request.sheet_url, request.page_content_column)
     documents = faq.similarity_search(vectordb, request.question, k=request.k)
     df_doc = util.transform_documents_to_dataframe(documents)
-    df_filter = util.remove_duplicates_by_column(df_doc, "ID")
+    df_filter = util.remove_duplicates_by_column(df_doc, request.id_column)
     return util.dataframe_to_dict(df_filter)
 
 

faq.py

@@ -103,7 +103,7 @@ def create_vectordb_id(
     if embedding_function is None:
         embedding_function = define_embedding_function(EMBEDDING_MODEL)
 
-    df = util.read_df(util.xlsx_url(faq_id))
+    df = util.read_df(util.xlsx_url(faq_id), page_content_column)
     documents = create_documents(df, page_content_column)
     vectordb = get_vectordb(
         faq_id=faq_id, embedding_function=embedding_function, documents=documents

util.py

@@ -4,6 +4,7 @@ SHEET_URL_X = "https://docs.google.com/spreadsheets/d/"
 SHEET_URL_Y = "/edit#gid="
 SHEET_URL_Y_EXPORT = "/export?gid="
 SPLIT_PAGE_BREAKS = False
+SYNONYMS = None
 
 
 def get_id(sheet_url: str) -> str:
@@ -17,10 +18,12 @@ def xlsx_url(get_id: str) -> str:
     return SHEET_URL_X + get_id[0:y] + SHEET_URL_Y_EXPORT + get_id[y + 1 :]
 
 
-def read_df(xlsx_url: str, split_page_breaks: bool = SPLIT_PAGE_BREAKS) -> pd.DataFrame:
+def read_df(xlsx_url: str, page_content_column: str) -> pd.DataFrame:
     df = pd.read_excel(xlsx_url, header=0, keep_default_na=False)
-    if split_page_breaks:
+    if SPLIT_PAGE_BREAKS:
         df = split_page_breaks(df, page_content_column)
+    if SYNONYMS is not None:
+        df = duplicate_rows_with_synonyms(df, page_content_column, SYNONYMS)
     return df
 
 
@@ -71,3 +74,20 @@ def dataframe_to_dict(df):
     df_records = df.to_dict(orient="records")
 
     return df_records
+
+
+def duplicate_rows_with_synonyms(df, column, synonyms):
+    new_rows = []
+    for index, row in df.iterrows():
+        new_rows.append(row)
+        for synonym_list in synonyms:
+            for word in row[column].split():
+                if word in synonym_list:
+                    for synonym in synonym_list:
+                        if synonym != word:
+                            new_row = row.copy()
+                            new_row[column] = row[column].replace(word, synonym)
+                            new_rows.append(new_row)
+    new_df = pd.DataFrame(new_rows, columns=df.columns)
+    new_df = new_df.reset_index(drop=True)
+    return new_df