Minor fixes

kstathou · Dec 5, 2023 · 2c88fb9 · 2c88fb9
1 parent 5751859
commit 2c88fb9
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/src/llm_stack/scripts/build_dataset_ner.py b/src/llm_stack/scripts/build_dataset_ner.py
@@ -13,7 +13,6 @@
 # Imports shared across functions
 with image.run_inside():
     import asyncio
-    import logging
     import os
 
     import pandas as pd
@@ -98,7 +97,6 @@ async def annotate_dataset_with_open_ai(
         messages = [
             system_message.to_prompt(),
             user_message.to_prompt(text=tup.abstract),
-            # formatting_message,
         ]
         tasks.append(
             openai_llm.generate(
@@ -146,8 +144,8 @@ async def main(
     local_data_path_ner_openai: str = "preprints_openai_ner.parquet",
 ) -> None:
     """Build an NER dataset using arXiv's papers and OpenAI's LLMs."""
-    logging.info("Fetching arXiv data...")
+    # Fetching the arXiv data
     fetch_arxiv_data.remote(local_data_path=local_data_path_raw)
 
-    logging.info("NER with OpenAI...")
+    # NER with OpenAI
     annotate_dataset_with_open_ai.remote(local_data_path=local_data_path_ner_openai)