Checked type hints

src/data_synthesis/generate_data.py

@@ -82,7 +82,7 @@ Please format the response as json in this style:
 
 
     @staticmethod
-    def reviews_for_product(product: Product, k: int):
+    def reviews_for_product(product: Product, k: int) -> str:
         prompt = f"Suggest exactly {k} reviews for this product.\nThe product is a {product.category.lower()[0:-1]} named the '{product.name}', which features {DataPrompt.format_features(product.features)}.\nFirst pick an integer star rating from 1 to 5 stars, where 1 is bad and 5 is great, for the review.\nNext write the review text of between 50 and 100 words for the review from the user. The text in the review should align to the star rating, so if the rating is 1 the review would be critical and if the rating is 5 the review would be positive.\n"
         prompt += """
 Please format the response as json in this style:
@@ -97,9 +97,10 @@ Please format the response as json in this style:
         return prompt
 
 
-def generate_products(category: str, features: List[str], k: int = 20):
+def generate_products(category: str, features: List[str], k: int = 20) -> None:
     """
     Call GPT3.5 Turbo model and get it to generate some products based on a category
+    Insert those products into the category
     """
     prompt = DataPrompt.products_for_category(category, features, k)
     response = openai.ChatCompletion.create(
@@ -214,7 +215,7 @@ def get_categories_and_features() -> Dict[str, List[str]]:
     return cats_and_feats
 
 
-def generate_all_products(target_count=40):
+def generate_all_products(target_count=40) -> None:
     """
     Generate all products for all categories, trying to reach a given target count
     of products.
@@ -235,7 +236,7 @@ def generate_all_products(target_count=40):
                 print(f"Skipping {cat_name} as targetting {target_count} and already have {len(existing_products)}")
 
 
-def dump_products_to_csv():
+def dump_products_to_csv() -> None:
     """
     Dump a csv file for debug, for every product showing category name and product name
     """
@@ -249,7 +250,7 @@ def dump_products_to_csv():
         f.write('\n'.join(cat_keys))
 
 
-def generate_reviews(target_count: int):
+def generate_reviews(target_count: int) -> None:
     """
     Generate reviews for each category up to a target count of reviews
     """
@@ -257,7 +258,7 @@ def generate_reviews(target_count: int):
         generate_reviews_for_category(cat, target_count)
 
 
-def generate_reviews_for_category(category: str, target_count: int):
+def generate_reviews_for_category(category: str, target_count: int) -> None:
     """
     Generate reviews for a specific category up to a given target number of reviews
     """
@@ -291,7 +292,7 @@ def generate_reviews_for_category(category: str, target_count: int):
                 print(f'{prod.category[:-1]}: {prod.name} has {len(prod.reviews)} reviews ({target_count} requested). Skipping.')
 
 
-def generate_reviews_for_product(product: Product, k: int):
+def generate_reviews_for_product(product: Product, k: int) -> None:
     """
     Generate a number of reviews from GPT3.5 for a specific product and add them to the product
     """
@@ -308,7 +309,7 @@ def generate_reviews_for_product(product: Product, k: int):
     add_reviews_to_product(output_text, product)
 
 
-def add_reviews_to_product(reviews_json: str, product: Product):
+def add_reviews_to_product(reviews_json: str, product: Product) -> None:
     """
     Load the reviews file containing this product category, append this review to the list and
     re-save the file