Fix schema access to title (#14)

Signed-off-by: Aivin V. Solatorio <avsolatorio@gmail.com>
worldbank · Jun 16, 2023 · 2641e6e · 2641e6e
1 parent 8f088e4
commit 2641e6e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/llm4data/schema/schema2info.py b/llm4data/schema/schema2info.py
@@ -3,7 +3,7 @@
 
 def get_doc_title(metadata: dict) -> str:
     """Get the title of the document from the metadata document_description."""
-    return metadata["title_statement"]["title"]
+    return metadata["document_description"]["title_statement"]["title"]
 
 
 def get_doc_authors(metadata: dict) -> list:

diff --git a/llm4data/scripts/indexing/docs/docs.py b/llm4data/scripts/indexing/docs/docs.py
@@ -11,6 +11,7 @@
 from llm4data.embeddings.docs import get_docs_embeddings
 from llm4data import index
 from llm4data import configs
+from llm4data.schema.schema2info import get_doc_title
 
 # Get the docs embeddings
 docs_embeddings = get_docs_embeddings()
@@ -39,7 +40,7 @@ def add_pdf_document(path: Union[str, Path], metadata: Optional[dict] = None):
         if len(documents):
             # Index the title of the document
             documents.append(
-                Document(page_content=metadata["title"], metadata=documents[0].metadata)
+                Document(page_content=get_doc_title(metadata), metadata=documents[0].metadata)
             )
 
         for doc in documents: