Merge pull request #6 from bespokelabsai/mahesh/fix

Minor refactoring and cleanups.

bella.py

@@ -1,15 +1,15 @@
 """Bella: Bespoke Labs Synthetic Data Generation Library."""
 
 import asyncio
-from typing import Optional
-from prompt import Prompter
-from datasets import Dataset
-import logging
 import json
+import logging
 import os
-from datasets.arrow_writer import ArrowWriter
+from typing import Optional
+
 from api_request_parallel_processor import process_api_requests_from_file
-import tiktoken
+from datasets import Dataset
+from datasets.arrow_writer import ArrowWriter
+from prompt import Prompter
 
 
 def _create_requests_file(

prompt.py

@@ -1,4 +1,3 @@
-from re import M
 from jinja2 import Template
 from pydantic import BaseModel
 from typing import Any, Dict, Optional, Type
@@ -9,8 +8,8 @@ class Prompter:
 
     def __init__(
         self,
-        model_name,
-        user_prompt,
+        model_name: str,
+        user_prompt: str,
         system_prompt: Optional[str] = None,
         response_format: Optional[Type[BaseModel]] = None,
     ):
@@ -32,14 +31,17 @@ def get_request_object(self, row: Dict[str, Any], idx: int) -> Dict[str, Any]:
         messages.append({"role": "user", "content": user_template.render(**row)})
 
         if self.response_format:
-            # OpenAI API https://platform.openai.com/docs/api-reference/chat/create#chat-create-response_format
+            # OpenAI API
+            # https://platform.openai.com/docs/api-reference/chat/create#chat-create-response_format
             request = {
                 "model": self.model_name,
                 "messages": messages,
                 "response_format": {
                     "type": "json_schema",
                     "json_schema": {
-                        "name": "output_schema",  # not sure if this should be something else. Also not sure if we should use strict: True
+                        # TODO(ryan): not sure if this should be something else.
+                        # TODO(ryan): also not sure if we should use strict: True
+                        "name": "output_schema",
                         "schema": self.response_format.model_json_schema(),
                     },
                 },