chore: prevent out of expected limit for scraped data size

Amber-Williams · Jan 1, 2024 · c16cde1 · c16cde1
1 parent d1764f4
commit c16cde1
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/main.py b/main.py
@@ -15,6 +15,9 @@
 
 batch_size = df['token_count'].sum() / settings.TOKEN_LIMIT
 batched_df = np.array_split(df, batch_size)
+if len(batched_df) > 30:
+    raise ValueError(f"Batch size is too large: {len(batched_df)}")
+
 chat = ChatExtractor(
         model_key=settings.OPENAI_API_KEY,
         model=settings.OPENAI_MODEL,