Full dataset generation

generate_synthetic_dataset.py

@@ -1,3 +1,5 @@
+import time
+
 from grazie.api.client.chat.prompt import ChatPrompt
 from grazie.api.client.endpoints import GrazieApiGatewayUrls
 from grazie.api.client.gateway import GrazieApiGatewayClient, GrazieAgent, AuthType
@@ -91,16 +93,17 @@ def generate_synthetic_dataset():
     df['initial_msg_prompt'] = df.apply(generate_prompt_for_row, axis=1)
     initial_messages_pred = []
 
-    for i, prompt in enumerate(tqdm(df['initial_msg_prompt'])):
+    for i, prompt in enumerate(tqdm(df['initial_msg_ prompt'])):
         output = None
 
-        if i < 5:
-            while output is None:
-                try:
-                    output = generate_initial_msg(prompt)
-                except:
-                    pass
-        initial_messages_pred.append(output if output is not None else "TBA")
+        while output is None:
+            try:
+                output = generate_initial_msg(prompt)
+            except:
+                time.sleep(0.5)
+
+        assert output is not None
+        initial_messages_pred.append(output)
 
     df['initial_msg_pred'] = initial_messages_pred