We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[Bug]: Send only needed columns in transforms
**Describe the bug**
Currently, in the webpage transform, even if a subset of columns are specified in the output columns, all columns are sent to the final dataset object.

**To Reproduce**
Run the example_webpage_transform and look at the ds.df

**Expected behavior**
We should just return the columns which are present in the output_columns dictionary.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/autolabel/transforms/webpage_transform.py`
Content:
```
1 from autolabel.schema import TransformType
2 from autolabel.transforms import BaseTransform
3 from typing import Dict, Any
4 import asyncio
5 import logging
6 import pandas as pd
7 import ssl
8 
9 from autolabel.cache import BaseCache
10 
11 logger = logging.getLogger(__name__)
12 
13 MAX_RETRIES = 3
14 MAX_KEEPALIVE_CONNECTIONS = 20
15 MAX_CONNECTIONS = 100
16 BACKOFF = 2
17 HEADERS = {}
18 HTML_PARSER = "html.parser"
19 
20 
21 class WebpageTransform(BaseTransform):
22     def __init__(
23         self,
24         cache: BaseCache,
25         output_columns: Dict[str, Any],
26         url_column: str,
27         timeout: int = 5,
28         headers: Dict[str, str] = HEADERS,
29     ) -> None:
30         super().__init__(cache, output_columns)
31         self.url_column = url_column
32         self.headers = headers
33         self.max_retries = MAX_RETRIES
34         try:
35             from bs4 import BeautifulSoup
36             import httpx
37 
38             if not headers.get("User-Agent"):
39                 from fake_useragent import UserAgent
40 
41                 headers["User-Agent"] = UserAgent().random
42 
43             self.httpx = httpx
44             self.timeout_time = timeout
45             self.timeout = httpx.Timeout(timeout)
46             limits = httpx.Limits(
47                 max_keepalive_connections=MAX_KEEPALIVE_CONNECTIONS,
48                 max_connections=MAX_CONNECTIONS,
49                 keepalive_expiry=timeout,
50             )
51             self.client = httpx.AsyncClient(
52                 timeout=self.timeout, limits=limits, follow_redirects=True
53             )
54             self.client_with_no_verify = httpx.AsyncClient(
55                 timeout=self.timeout, limits=limits, follow_redirects=True, verify=False
56             )
57             self.beautiful_soup = BeautifulSoup
58         except ImportError:
59             raise ImportError(
60                 "BeautifulSoup, httpx and fake_useragent are required to use the webpage transform. Please install them with the following command: pip install beautifulsoup4 httpx fake_useragent"
61             )
62 
63     def name(self) -> str:
64         return TransformType.WEBPAGE_TRANSFORM
65 
66     @property
67     def output_columns(self) -> Dict[str, Any]:
68         COLUMN_NAMES = [
69             "content_column",
70             "content_in_bytes_column",
71             "soup_column",
72             "metadata_column",
73         ]
74         return {k: self._output_columns.get(k, k) for k in COLUMN_NAMES}
75 
76     def _load_metadata(self, url, soup) -> Dict[str, Any]:
77         metadata = {"url": url}
78         if soup.find("title"):
79             metadata["title"] = soup.find("title").get_text()
80         for meta in soup.find_all("meta"):
81             if meta.get("name") and meta.get("content"):
82                 metadata[meta.get("name")] = meta.get("content")
83             elif meta.get("property") and meta.get("content"):
84                 metadata[meta.get("property")] = meta.get("content")
85         return metadata
86 
87     async def _load_url(
88         self, url: str, verify=True, headers=HEADERS, retry_count=0
89     ) -> Dict[str, Any]:
90         if retry_count >= self.max_retries:
91             logger.warning(f"Max retries reached for URL: {url}")
92             return {}
93 
94         try:
95             client = self.client
96             if not verify:
97                 client = self.client_with_no_verify
98             response = await client.get(url, headers=headers)
99 
100             # TODO: Add support for other parsers
101             content_bytes = response.content
102             soup = self.beautiful_soup(content_bytes, HTML_PARSER)
103             return {
104                 "content": soup.get_text(),
105                 "content_bytes": content_bytes,
106                 "soup": soup,
107                 "metadata": self._load_metadata(url, soup),
108             }
109         except self.httpx.ConnectTimeout as e:
110             logger.error(f"Timeout when fetching content from URL: {url}")
111             return {}
112         except ssl.SSLCertVerificationError as e:
113             logger.warning(
114                 f"SSL verification error when fetching content from URL: {url}, retrying with verify=False"
115             )
116             await asyncio.sleep(BACKOFF**retry_count)
117             return await self._load_url(
118                 url, verify=False, headers=headers, retry_count=retry_count + 1
119             )
120         except Exception as e:
121             logger.error(f"Error fetching content from URL: {url}. Exception: {e}")
122             return {}
123 
124     async def _apply(self, row: Dict[str, Any]) -> Dict[str, Any]:
125         url = row[self.url_column]
126         url_response_data = {}
127         if pd.isna(url):
128             logger.warning(f"Empty url in row {row}")
129         else:
130             url_response_data = await self._load_url(url)
131 
132         transformed_row = {
133             self.output_columns["content_column"]: url_response_data.get("content"),
134             self.output_columns["content_in_bytes_column"]: url_response_data.get(
135                 "content_bytes"
136             ),
137             self.output_columns["soup_column"]: url_response_data.get("soup"),
138             self.output_columns["metadata_column"]: url_response_data.get("metadata"),
139         }
140 
141         return transformed_row
142 
143     def params(self):
144         return {
145             "url_column": self.url_column,
146             "output_columns": self.output_columns,
147             "timeout": self.timeout_time,
148         }
149
```

Path: `src/autolabel/transforms/pdf.py`
Content:
```
1 from typing import List, Dict, Any
2 
3 from autolabel.schema import TransformType
4 from autolabel.transforms import BaseTransform
5 from autolabel.cache import BaseCache
6 
7 
8 class PDFTransform(BaseTransform):
9     def __init__(
10         self,
11         cache: BaseCache,
12         output_columns: Dict[str, Any],
13         file_path_column: str,
14         ocr_enabled: bool = False,
15         page_header: str = "Page {page_num}: {page_content}",
16         page_sep: str = "\n\n",
17     ) -> None:
18         """The output columns for this class should be in the order: [content_column, num_pages_column]"""
19         super().__init__(cache, output_columns)
20         self.file_path_column = file_path_column
21         self.ocr_enabled = ocr_enabled
22         self.page_format = page_header
23         self.page_sep = page_sep
24 
25         if self.ocr_enabled:
26             try:
27                 from pdf2image import convert_from_path
28                 import pytesseract
29 
30                 self.convert_from_path = convert_from_path
31                 self.pytesseract = pytesseract
32                 self.pytesseract.get_tesseract_version()
33             except ImportError:
34                 raise ImportError(
35                     "pdf2image and pytesseract are required to use the pdf transform with ocr. Please install pdf2image and pytesseract with the following command: pip install pdf2image pytesseract"
36                 )
37             except EnvironmentError:
38                 raise EnvironmentError(
39                     "The tesseract engine is required to use the pdf transform with ocr. Please see https://tesseract-ocr.github.io/tessdoc/Installation.html for installation instructions."
40                 )
41         else:
42             try:
43                 from langchain.document_loaders import PDFPlumberLoader
44 
45                 self.PDFPlumberLoader = PDFPlumberLoader
46             except ImportError:
47                 raise ImportError(
48                     "pdfplumber is required to use the pdf transform. Please install pdfplumber with the following command: pip install pdfplumber"
49                 )
50 
51     @staticmethod
52     def name() -> str:
53         return TransformType.PDF
54 
55     @property
56     def output_columns(self) -> Dict[str, Any]:
57         COLUMN_NAMES = [
58             "content_column",
59             "metadata_column",
60         ]
61         return {k: self._output_columns.get(k, k) for k in COLUMN_NAMES}
62 
63     def get_page_texts(self, row: Dict[str, Any]) -> List[str]:
64         """This function gets the text from each page of a PDF file.
65         If OCR is enabled, it uses the pdf2image library to convert the PDF into images and then uses
66         pytesseract to convert the images into text. Otherwise, it uses pdfplumber to extract the text.
67 
68         Args:
69             row (Dict[str, Any]): The row of data to be transformed.
70 
71         Returns:
72             List[str]: A list of strings containing the text from each page of the PDF.
73         """
74         if self.ocr_enabled:
75             pages = self.convert_from_path(row[self.file_path_column])
76             return [self.pytesseract.image_to_string(page) for page in pages]
77         else:
78             loader = self.PDFPlumberLoader(row[self.file_path_column])
79             return [page.page_content for page in loader.load()]
80 
81     async def _apply(self, row: Dict[str, Any]) -> Dict[str, Any]:
82         """This function transforms a PDF file into a string of text.
83         The text is formatted according to the page_format and
84         page_sep parameters and returned as a string.
85 
86         Args:
87             row (Dict[str, Any]): The row of data to be transformed.
88 
89         Returns:
90             Dict[str, Any]: The dict of output columns.
91         """
92         texts = []
93         for idx, text in enumerate(self.get_page_texts(row)):
94             texts.append(self.page_format.format(page_num=idx + 1, page_content=text))
95         output = self.page_sep.join(texts)
96         transformed_row = {
97             self.output_columns["content_column"]: output,
98             self.output_columns["metadata_column"]: {"num_pages": len(texts)},
99         }
100         return transformed_row
101 
102     def params(self):
103         return {
104             "file_path_column": self.file_path_column,
105             "ocr_enabled": self.ocr_enabled,
106             "page_header": self.page_format,
107             "page_sep": self.page_sep,
108             "output_columns": self.output_columns,
109         }
110
```

Path: `src/autolabel/transforms/base.py`
Content:
```
1 from abc import ABC, abstractmethod
2 from typing import Dict, Any
3 from autolabel.cache import BaseCache
4 from autolabel.schema import TransformCacheEntry
5 
6 
7 class BaseTransform(ABC):
8     """Base class for all transforms."""
9 
10     TTL_MS = 60 * 60 * 24 * 7 * 1000  # 1 week
11 
12     def __init__(self, cache: BaseCache, output_columns: Dict[str, Any]) -> None:
13         """
14         Initialize a transform.
15         Args:
16             cache: A cache object to use for caching the results of this transform.
17             output_columns: A dictionary of output columns. The keys are the names of the output columns as expected by the transform. The values are the column names they should be mapped to in the dataset.
18         """
19         super().__init__()
20         self._output_columns = output_columns
21         self.cache = cache
22 
23     @staticmethod
24     @abstractmethod
25     def name() -> str:
26         """
27         Returns the name of the transform.
28         """
29         pass
30 
31     @property
32     @abstractmethod
33     def output_columns(self) -> Dict[str, Any]:
34         """
35         Returns a dictionary of output columns. The keys are the names of the output columns
36         as expected by the transform. The values are the column names they should be mapped to in
37         the dataset.
38         """
39         return self._output_columns
40 
41     @abstractmethod
42     async def _apply(self, row: Dict[str, Any]) -> Dict[str, Any]:
43         """
44         Applies the transform to the given row.
45         Args:
46             row: A dictionary representing a row in the dataset. The keys are the column names and the values are the column values.
47         Returns:
48             A dictionary representing the transformed row. The keys are the column names and the values are the column values.
49         """
50         pass
51 
52     @abstractmethod
53     def params(self) -> Dict[str, Any]:
54         """
55         Returns a dictionary of parameters that can be used to uniquely identify this transform.
56         Returns:
57             A dictionary of parameters that can be used to uniquely identify this transform.
58         """
59         return {}
60 
61     async def apply(self, row: Dict[str, Any]) -> Dict[str, Any]:
62         if self.cache is not None:
63             cache_entry = TransformCacheEntry(
64                 transform_name=self.name(),
65                 transform_params=self.params(),
66                 input=row,
67                 ttl_ms=self.TTL_MS,
68             )
69             output = self.cache.lookup(cache_entry)
70 
71             if output is not None:
72                 # Cache hit
73                 return output
74 
75         output = await self._apply(row)
76 
77         if self.cache is not None:
78             cache_entry.output = output
79             self.cache.update(cache_entry)
80 
81         return output
82
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/src/autolabel/transforms/base.py b/src/autolabel/transforms/base.py
--- a/src/autolabel/transforms/base.py
+++ b/src/autolabel/transforms/base.py
@@ -29,14 +29,13 @@
         pass
 
     @property
-    @abstractmethod
     def output_columns(self) -> Dict[str, Any]:
         """
         Returns a dictionary of output columns. The keys are the names of the output columns
         as expected by the transform. The values are the column names they should be mapped to in
         the dataset.
         """
-        return self._output_columns
+        return {k: self._output_columns.get(k, None) for k in self.COLUMN_NAMES}
 
     @abstractmethod
     async def _apply(self, row: Dict[str, Any]) -> Dict[str, Any]:
@@ -79,3 +78,15 @@
             self.cache.update(cache_entry)
 
         return output
+
+    def _return_output_row(self, row: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Returns the output row with the correct column names.
+        Args:
+            row: The output row.
+        Returns:
+            The output row with the correct column names.
+        """
+        # remove null key
+        row.pop(None, None)
+        return row
diff --git a/src/autolabel/transforms/pdf.py b/src/autolabel/transforms/pdf.py
--- a/src/autolabel/transforms/pdf.py
+++ b/src/autolabel/transforms/pdf.py
@@ -6,6 +6,11 @@
 
 
 class PDFTransform(BaseTransform):
+    COLUMN_NAMES = [
+        "content_column",
+        "metadata_column",
+    ]
+
     def __init__(
         self,
         cache: BaseCache,
@@ -52,14 +57,6 @@
     def name() -> str:
         return TransformType.PDF
 
-    @property
-    def output_columns(self) -> Dict[str, Any]:
-        COLUMN_NAMES = [
-            "content_column",
-            "metadata_column",
-        ]
-        return {k: self._output_columns.get(k, k) for k in COLUMN_NAMES}
-
     def get_page_texts(self, row: Dict[str, Any]) -> List[str]:
         """This function gets the text from each page of a PDF file.
         If OCR is enabled, it uses the pdf2image library to convert the PDF into images and then uses
@@ -97,7 +94,7 @@
             self.output_columns["content_column"]: output,
             self.output_columns["metadata_column"]: {"num_pages": len(texts)},
         }
-        return transformed_row
+        return self._return_output_row(transformed_row)
 
     def params(self):
         return {
diff --git a/src/autolabel/transforms/webpage_transform.py b/src/autolabel/transforms/webpage_transform.py
--- a/src/autolabel/transforms/webpage_transform.py
+++ b/src/autolabel/transforms/webpage_transform.py
@@ -19,6 +19,13 @@
 
 
 class WebpageTransform(BaseTransform):
+    COLUMN_NAMES = [
+        "content_column",
+        "content_in_bytes_column",
+        "soup_column",
+        "metadata_column",
+    ]
+
     def __init__(
         self,
         cache: BaseCache,
@@ -63,16 +70,6 @@
     def name(self) -> str:
         return TransformType.WEBPAGE_TRANSFORM
 
-    @property
-    def output_columns(self) -> Dict[str, Any]:
-        COLUMN_NAMES = [
-            "content_column",
-            "content_in_bytes_column",
-            "soup_column",
-            "metadata_column",
-        ]
-        return {k: self._output_columns.get(k, k) for k in COLUMN_NAMES}
-
     def _load_metadata(self, url, soup) -> Dict[str, Any]:
         metadata = {"url": url}
         if soup.find("title"):
@@ -138,7 +135,7 @@
             self.output_columns["metadata_column"]: url_response_data.get("metadata"),
         }
 
-        return transformed_row
+        return self._return_output_row(transformed_row)
 
     def params(self):
         return {

{"golden_diff": "diff --git a/src/autolabel/transforms/base.py b/src/autolabel/transforms/base.py\n--- a/src/autolabel/transforms/base.py\n+++ b/src/autolabel/transforms/base.py\n@@ -29,14 +29,13 @@\n         pass\n \n     @property\n-    @abstractmethod\n     def output_columns(self) -> Dict[str, Any]:\n         \"\"\"\n         Returns a dictionary of output columns. The keys are the names of the output columns\n         as expected by the transform. The values are the column names they should be mapped to in\n         the dataset.\n         \"\"\"\n-        return self._output_columns\n+        return {k: self._output_columns.get(k, None) for k in self.COLUMN_NAMES}\n \n     @abstractmethod\n     async def _apply(self, row: Dict[str, Any]) -> Dict[str, Any]:\n@@ -79,3 +78,15 @@\n             self.cache.update(cache_entry)\n \n         return output\n+\n+    def _return_output_row(self, row: Dict[str, Any]) -> Dict[str, Any]:\n+        \"\"\"\n+        Returns the output row with the correct column names.\n+        Args:\n+            row: The output row.\n+        Returns:\n+            The output row with the correct column names.\n+        \"\"\"\n+        # remove null key\n+        row.pop(None, None)\n+        return row\ndiff --git a/src/autolabel/transforms/pdf.py b/src/autolabel/transforms/pdf.py\n--- a/src/autolabel/transforms/pdf.py\n+++ b/src/autolabel/transforms/pdf.py\n@@ -6,6 +6,11 @@\n \n \n class PDFTransform(BaseTransform):\n+    COLUMN_NAMES = [\n+        \"content_column\",\n+        \"metadata_column\",\n+    ]\n+\n     def __init__(\n         self,\n         cache: BaseCache,\n@@ -52,14 +57,6 @@\n     def name() -> str:\n         return TransformType.PDF\n \n-    @property\n-    def output_columns(self) -> Dict[str, Any]:\n-        COLUMN_NAMES = [\n-            \"content_column\",\n-            \"metadata_column\",\n-        ]\n-        return {k: self._output_columns.get(k, k) for k in COLUMN_NAMES}\n-\n     def get_page_texts(self, row: Dict[str, Any]) -> List[str]:\n         \"\"\"This function gets the text from each page of a PDF file.\n         If OCR is enabled, it uses the pdf2image library to convert the PDF into images and then uses\n@@ -97,7 +94,7 @@\n             self.output_columns[\"content_column\"]: output,\n             self.output_columns[\"metadata_column\"]: {\"num_pages\": len(texts)},\n         }\n-        return transformed_row\n+        return self._return_output_row(transformed_row)\n \n     def params(self):\n         return {\ndiff --git a/src/autolabel/transforms/webpage_transform.py b/src/autolabel/transforms/webpage_transform.py\n--- a/src/autolabel/transforms/webpage_transform.py\n+++ b/src/autolabel/transforms/webpage_transform.py\n@@ -19,6 +19,13 @@\n \n \n class WebpageTransform(BaseTransform):\n+    COLUMN_NAMES = [\n+        \"content_column\",\n+        \"content_in_bytes_column\",\n+        \"soup_column\",\n+        \"metadata_column\",\n+    ]\n+\n     def __init__(\n         self,\n         cache: BaseCache,\n@@ -63,16 +70,6 @@\n     def name(self) -> str:\n         return TransformType.WEBPAGE_TRANSFORM\n \n-    @property\n-    def output_columns(self) -> Dict[str, Any]:\n-        COLUMN_NAMES = [\n-            \"content_column\",\n-            \"content_in_bytes_column\",\n-            \"soup_column\",\n-            \"metadata_column\",\n-        ]\n-        return {k: self._output_columns.get(k, k) for k in COLUMN_NAMES}\n-\n     def _load_metadata(self, url, soup) -> Dict[str, Any]:\n         metadata = {\"url\": url}\n         if soup.find(\"title\"):\n@@ -138,7 +135,7 @@\n             self.output_columns[\"metadata_column\"]: url_response_data.get(\"metadata\"),\n         }\n \n-        return transformed_row\n+        return self._return_output_row(transformed_row)\n \n     def params(self):\n         return {\n", "issue": "[Bug]: Send only needed columns in transforms\n**Describe the bug**\r\nCurrently, in the webpage transform, even if a subset of columns are specified in the output columns, all columns are sent to the final dataset object.\r\n\r\n**To Reproduce**\r\nRun the example_webpage_transform and look at the ds.df\r\n\r\n**Expected behavior**\r\nWe should just return the columns which are present in the output_columns dictionary.\r\n\n", "before_files": [{"content": "from autolabel.schema import TransformType\nfrom autolabel.transforms import BaseTransform\nfrom typing import Dict, Any\nimport asyncio\nimport logging\nimport pandas as pd\nimport ssl\n\nfrom autolabel.cache import BaseCache\n\nlogger = logging.getLogger(__name__)\n\nMAX_RETRIES = 3\nMAX_KEEPALIVE_CONNECTIONS = 20\nMAX_CONNECTIONS = 100\nBACKOFF = 2\nHEADERS = {}\nHTML_PARSER = \"html.parser\"\n\n\nclass WebpageTransform(BaseTransform):\n    def __init__(\n        self,\n        cache: BaseCache,\n        output_columns: Dict[str, Any],\n        url_column: str,\n        timeout: int = 5,\n        headers: Dict[str, str] = HEADERS,\n    ) -> None:\n        super().__init__(cache, output_columns)\n        self.url_column = url_column\n        self.headers = headers\n        self.max_retries = MAX_RETRIES\n        try:\n            from bs4 import BeautifulSoup\n            import httpx\n\n            if not headers.get(\"User-Agent\"):\n                from fake_useragent import UserAgent\n\n                headers[\"User-Agent\"] = UserAgent().random\n\n            self.httpx = httpx\n            self.timeout_time = timeout\n            self.timeout = httpx.Timeout(timeout)\n            limits = httpx.Limits(\n                max_keepalive_connections=MAX_KEEPALIVE_CONNECTIONS,\n                max_connections=MAX_CONNECTIONS,\n                keepalive_expiry=timeout,\n            )\n            self.client = httpx.AsyncClient(\n                timeout=self.timeout, limits=limits, follow_redirects=True\n            )\n            self.client_with_no_verify = httpx.AsyncClient(\n                timeout=self.timeout, limits=limits, follow_redirects=True, verify=False\n            )\n            self.beautiful_soup = BeautifulSoup\n        except ImportError:\n            raise ImportError(\n                \"BeautifulSoup, httpx and fake_useragent are required to use the webpage transform. Please install them with the following command: pip install beautifulsoup4 httpx fake_useragent\"\n            )\n\n    def name(self) -> str:\n        return TransformType.WEBPAGE_TRANSFORM\n\n    @property\n    def output_columns(self) -> Dict[str, Any]:\n        COLUMN_NAMES = [\n            \"content_column\",\n            \"content_in_bytes_column\",\n            \"soup_column\",\n            \"metadata_column\",\n        ]\n        return {k: self._output_columns.get(k, k) for k in COLUMN_NAMES}\n\n    def _load_metadata(self, url, soup) -> Dict[str, Any]:\n        metadata = {\"url\": url}\n        if soup.find(\"title\"):\n            metadata[\"title\"] = soup.find(\"title\").get_text()\n        for meta in soup.find_all(\"meta\"):\n            if meta.get(\"name\") and meta.get(\"content\"):\n                metadata[meta.get(\"name\")] = meta.get(\"content\")\n            elif meta.get(\"property\") and meta.get(\"content\"):\n                metadata[meta.get(\"property\")] = meta.get(\"content\")\n        return metadata\n\n    async def _load_url(\n        self, url: str, verify=True, headers=HEADERS, retry_count=0\n    ) -> Dict[str, Any]:\n        if retry_count >= self.max_retries:\n            logger.warning(f\"Max retries reached for URL: {url}\")\n            return {}\n\n        try:\n            client = self.client\n            if not verify:\n                client = self.client_with_no_verify\n            response = await client.get(url, headers=headers)\n\n            # TODO: Add support for other parsers\n            content_bytes = response.content\n            soup = self.beautiful_soup(content_bytes, HTML_PARSER)\n            return {\n                \"content\": soup.get_text(),\n                \"content_bytes\": content_bytes,\n                \"soup\": soup,\n                \"metadata\": self._load_metadata(url, soup),\n            }\n        except self.httpx.ConnectTimeout as e:\n            logger.error(f\"Timeout when fetching content from URL: {url}\")\n            return {}\n        except ssl.SSLCertVerificationError as e:\n            logger.warning(\n                f\"SSL verification error when fetching content from URL: {url}, retrying with verify=False\"\n            )\n            await asyncio.sleep(BACKOFF**retry_count)\n            return await self._load_url(\n                url, verify=False, headers=headers, retry_count=retry_count + 1\n            )\n        except Exception as e:\n            logger.error(f\"Error fetching content from URL: {url}. Exception: {e}\")\n            return {}\n\n    async def _apply(self, row: Dict[str, Any]) -> Dict[str, Any]:\n        url = row[self.url_column]\n        url_response_data = {}\n        if pd.isna(url):\n            logger.warning(f\"Empty url in row {row}\")\n        else:\n            url_response_data = await self._load_url(url)\n\n        transformed_row = {\n            self.output_columns[\"content_column\"]: url_response_data.get(\"content\"),\n            self.output_columns[\"content_in_bytes_column\"]: url_response_data.get(\n                \"content_bytes\"\n            ),\n            self.output_columns[\"soup_column\"]: url_response_data.get(\"soup\"),\n            self.output_columns[\"metadata_column\"]: url_response_data.get(\"metadata\"),\n        }\n\n        return transformed_row\n\n    def params(self):\n        return {\n            \"url_column\": self.url_column,\n            \"output_columns\": self.output_columns,\n            \"timeout\": self.timeout_time,\n        }\n", "path": "src/autolabel/transforms/webpage_transform.py"}, {"content": "from typing import List, Dict, Any\n\nfrom autolabel.schema import TransformType\nfrom autolabel.transforms import BaseTransform\nfrom autolabel.cache import BaseCache\n\n\nclass PDFTransform(BaseTransform):\n    def __init__(\n        self,\n        cache: BaseCache,\n        output_columns: Dict[str, Any],\n        file_path_column: str,\n        ocr_enabled: bool = False,\n        page_header: str = \"Page {page_num}: {page_content}\",\n        page_sep: str = \"\\n\\n\",\n    ) -> None:\n        \"\"\"The output columns for this class should be in the order: [content_column, num_pages_column]\"\"\"\n        super().__init__(cache, output_columns)\n        self.file_path_column = file_path_column\n        self.ocr_enabled = ocr_enabled\n        self.page_format = page_header\n        self.page_sep = page_sep\n\n        if self.ocr_enabled:\n            try:\n                from pdf2image import convert_from_path\n                import pytesseract\n\n                self.convert_from_path = convert_from_path\n                self.pytesseract = pytesseract\n                self.pytesseract.get_tesseract_version()\n            except ImportError:\n                raise ImportError(\n                    \"pdf2image and pytesseract are required to use the pdf transform with ocr. Please install pdf2image and pytesseract with the following command: pip install pdf2image pytesseract\"\n                )\n            except EnvironmentError:\n                raise EnvironmentError(\n                    \"The tesseract engine is required to use the pdf transform with ocr. Please see https://tesseract-ocr.github.io/tessdoc/Installation.html for installation instructions.\"\n                )\n        else:\n            try:\n                from langchain.document_loaders import PDFPlumberLoader\n\n                self.PDFPlumberLoader = PDFPlumberLoader\n            except ImportError:\n                raise ImportError(\n                    \"pdfplumber is required to use the pdf transform. Please install pdfplumber with the following command: pip install pdfplumber\"\n                )\n\n    @staticmethod\n    def name() -> str:\n        return TransformType.PDF\n\n    @property\n    def output_columns(self) -> Dict[str, Any]:\n        COLUMN_NAMES = [\n            \"content_column\",\n            \"metadata_column\",\n        ]\n        return {k: self._output_columns.get(k, k) for k in COLUMN_NAMES}\n\n    def get_page_texts(self, row: Dict[str, Any]) -> List[str]:\n        \"\"\"This function gets the text from each page of a PDF file.\n        If OCR is enabled, it uses the pdf2image library to convert the PDF into images and then uses\n        pytesseract to convert the images into text. Otherwise, it uses pdfplumber to extract the text.\n\n        Args:\n            row (Dict[str, Any]): The row of data to be transformed.\n\n        Returns:\n            List[str]: A list of strings containing the text from each page of the PDF.\n        \"\"\"\n        if self.ocr_enabled:\n            pages = self.convert_from_path(row[self.file_path_column])\n            return [self.pytesseract.image_to_string(page) for page in pages]\n        else:\n            loader = self.PDFPlumberLoader(row[self.file_path_column])\n            return [page.page_content for page in loader.load()]\n\n    async def _apply(self, row: Dict[str, Any]) -> Dict[str, Any]:\n        \"\"\"This function transforms a PDF file into a string of text.\n        The text is formatted according to the page_format and\n        page_sep parameters and returned as a string.\n\n        Args:\n            row (Dict[str, Any]): The row of data to be transformed.\n\n        Returns:\n            Dict[str, Any]: The dict of output columns.\n        \"\"\"\n        texts = []\n        for idx, text in enumerate(self.get_page_texts(row)):\n            texts.append(self.page_format.format(page_num=idx + 1, page_content=text))\n        output = self.page_sep.join(texts)\n        transformed_row = {\n            self.output_columns[\"content_column\"]: output,\n            self.output_columns[\"metadata_column\"]: {\"num_pages\": len(texts)},\n        }\n        return transformed_row\n\n    def params(self):\n        return {\n            \"file_path_column\": self.file_path_column,\n            \"ocr_enabled\": self.ocr_enabled,\n            \"page_header\": self.page_format,\n            \"page_sep\": self.page_sep,\n            \"output_columns\": self.output_columns,\n        }\n", "path": "src/autolabel/transforms/pdf.py"}, {"content": "from abc import ABC, abstractmethod\nfrom typing import Dict, Any\nfrom autolabel.cache import BaseCache\nfrom autolabel.schema import TransformCacheEntry\n\n\nclass BaseTransform(ABC):\n    \"\"\"Base class for all transforms.\"\"\"\n\n    TTL_MS = 60 * 60 * 24 * 7 * 1000  # 1 week\n\n    def __init__(self, cache: BaseCache, output_columns: Dict[str, Any]) -> None:\n        \"\"\"\n        Initialize a transform.\n        Args:\n            cache: A cache object to use for caching the results of this transform.\n            output_columns: A dictionary of output columns. The keys are the names of the output columns as expected by the transform. The values are the column names they should be mapped to in the dataset.\n        \"\"\"\n        super().__init__()\n        self._output_columns = output_columns\n        self.cache = cache\n\n    @staticmethod\n    @abstractmethod\n    def name() -> str:\n        \"\"\"\n        Returns the name of the transform.\n        \"\"\"\n        pass\n\n    @property\n    @abstractmethod\n    def output_columns(self) -> Dict[str, Any]:\n        \"\"\"\n        Returns a dictionary of output columns. The keys are the names of the output columns\n        as expected by the transform. The values are the column names they should be mapped to in\n        the dataset.\n        \"\"\"\n        return self._output_columns\n\n    @abstractmethod\n    async def _apply(self, row: Dict[str, Any]) -> Dict[str, Any]:\n        \"\"\"\n        Applies the transform to the given row.\n        Args:\n            row: A dictionary representing a row in the dataset. The keys are the column names and the values are the column values.\n        Returns:\n            A dictionary representing the transformed row. The keys are the column names and the values are the column values.\n        \"\"\"\n        pass\n\n    @abstractmethod\n    def params(self) -> Dict[str, Any]:\n        \"\"\"\n        Returns a dictionary of parameters that can be used to uniquely identify this transform.\n        Returns:\n            A dictionary of parameters that can be used to uniquely identify this transform.\n        \"\"\"\n        return {}\n\n    async def apply(self, row: Dict[str, Any]) -> Dict[str, Any]:\n        if self.cache is not None:\n            cache_entry = TransformCacheEntry(\n                transform_name=self.name(),\n                transform_params=self.params(),\n                input=row,\n                ttl_ms=self.TTL_MS,\n            )\n            output = self.cache.lookup(cache_entry)\n\n            if output is not None:\n                # Cache hit\n                return output\n\n        output = await self._apply(row)\n\n        if self.cache is not None:\n            cache_entry.output = output\n            self.cache.update(cache_entry)\n\n        return output\n", "path": "src/autolabel/transforms/base.py"}], "after_files": [{"content": "from autolabel.schema import TransformType\nfrom autolabel.transforms import BaseTransform\nfrom typing import Dict, Any\nimport asyncio\nimport logging\nimport pandas as pd\nimport ssl\n\nfrom autolabel.cache import BaseCache\n\nlogger = logging.getLogger(__name__)\n\nMAX_RETRIES = 3\nMAX_KEEPALIVE_CONNECTIONS = 20\nMAX_CONNECTIONS = 100\nBACKOFF = 2\nHEADERS = {}\nHTML_PARSER = \"html.parser\"\n\n\nclass WebpageTransform(BaseTransform):\n    COLUMN_NAMES = [\n        \"content_column\",\n        \"content_in_bytes_column\",\n        \"soup_column\",\n        \"metadata_column\",\n    ]\n\n    def __init__(\n        self,\n        cache: BaseCache,\n        output_columns: Dict[str, Any],\n        url_column: str,\n        timeout: int = 5,\n        headers: Dict[str, str] = HEADERS,\n    ) -> None:\n        super().__init__(cache, output_columns)\n        self.url_column = url_column\n        self.headers = headers\n        self.max_retries = MAX_RETRIES\n        try:\n            from bs4 import BeautifulSoup\n            import httpx\n\n            if not headers.get(\"User-Agent\"):\n                from fake_useragent import UserAgent\n\n                headers[\"User-Agent\"] = UserAgent().random\n\n            self.httpx = httpx\n            self.timeout_time = timeout\n            self.timeout = httpx.Timeout(timeout)\n            limits = httpx.Limits(\n                max_keepalive_connections=MAX_KEEPALIVE_CONNECTIONS,\n                max_connections=MAX_CONNECTIONS,\n                keepalive_expiry=timeout,\n            )\n            self.client = httpx.AsyncClient(\n                timeout=self.timeout, limits=limits, follow_redirects=True\n            )\n            self.client_with_no_verify = httpx.AsyncClient(\n                timeout=self.timeout, limits=limits, follow_redirects=True, verify=False\n            )\n            self.beautiful_soup = BeautifulSoup\n        except ImportError:\n            raise ImportError(\n                \"BeautifulSoup, httpx and fake_useragent are required to use the webpage transform. Please install them with the following command: pip install beautifulsoup4 httpx fake_useragent\"\n            )\n\n    def name(self) -> str:\n        return TransformType.WEBPAGE_TRANSFORM\n\n    def _load_metadata(self, url, soup) -> Dict[str, Any]:\n        metadata = {\"url\": url}\n        if soup.find(\"title\"):\n            metadata[\"title\"] = soup.find(\"title\").get_text()\n        for meta in soup.find_all(\"meta\"):\n            if meta.get(\"name\") and meta.get(\"content\"):\n                metadata[meta.get(\"name\")] = meta.get(\"content\")\n            elif meta.get(\"property\") and meta.get(\"content\"):\n                metadata[meta.get(\"property\")] = meta.get(\"content\")\n        return metadata\n\n    async def _load_url(\n        self, url: str, verify=True, headers=HEADERS, retry_count=0\n    ) -> Dict[str, Any]:\n        if retry_count >= self.max_retries:\n            logger.warning(f\"Max retries reached for URL: {url}\")\n            return {}\n\n        try:\n            client = self.client\n            if not verify:\n                client = self.client_with_no_verify\n            response = await client.get(url, headers=headers)\n\n            # TODO: Add support for other parsers\n            content_bytes = response.content\n            soup = self.beautiful_soup(content_bytes, HTML_PARSER)\n            return {\n                \"content\": soup.get_text(),\n                \"content_bytes\": content_bytes,\n                \"soup\": soup,\n                \"metadata\": self._load_metadata(url, soup),\n            }\n        except self.httpx.ConnectTimeout as e:\n            logger.error(f\"Timeout when fetching content from URL: {url}\")\n            return {}\n        except ssl.SSLCertVerificationError as e:\n            logger.warning(\n                f\"SSL verification error when fetching content from URL: {url}, retrying with verify=False\"\n            )\n            await asyncio.sleep(BACKOFF**retry_count)\n            return await self._load_url(\n                url, verify=False, headers=headers, retry_count=retry_count + 1\n            )\n        except Exception as e:\n            logger.error(f\"Error fetching content from URL: {url}. Exception: {e}\")\n            return {}\n\n    async def _apply(self, row: Dict[str, Any]) -> Dict[str, Any]:\n        url = row[self.url_column]\n        url_response_data = {}\n        if pd.isna(url):\n            logger.warning(f\"Empty url in row {row}\")\n        else:\n            url_response_data = await self._load_url(url)\n\n        transformed_row = {\n            self.output_columns[\"content_column\"]: url_response_data.get(\"content\"),\n            self.output_columns[\"content_in_bytes_column\"]: url_response_data.get(\n                \"content_bytes\"\n            ),\n            self.output_columns[\"soup_column\"]: url_response_data.get(\"soup\"),\n            self.output_columns[\"metadata_column\"]: url_response_data.get(\"metadata\"),\n        }\n\n        return self._return_output_row(transformed_row)\n\n    def params(self):\n        return {\n            \"url_column\": self.url_column,\n            \"output_columns\": self.output_columns,\n            \"timeout\": self.timeout_time,\n        }\n", "path": "src/autolabel/transforms/webpage_transform.py"}, {"content": "from typing import List, Dict, Any\n\nfrom autolabel.schema import TransformType\nfrom autolabel.transforms import BaseTransform\nfrom autolabel.cache import BaseCache\n\n\nclass PDFTransform(BaseTransform):\n    COLUMN_NAMES = [\n        \"content_column\",\n        \"metadata_column\",\n    ]\n\n    def __init__(\n        self,\n        cache: BaseCache,\n        output_columns: Dict[str, Any],\n        file_path_column: str,\n        ocr_enabled: bool = False,\n        page_header: str = \"Page {page_num}: {page_content}\",\n        page_sep: str = \"\\n\\n\",\n    ) -> None:\n        \"\"\"The output columns for this class should be in the order: [content_column, num_pages_column]\"\"\"\n        super().__init__(cache, output_columns)\n        self.file_path_column = file_path_column\n        self.ocr_enabled = ocr_enabled\n        self.page_format = page_header\n        self.page_sep = page_sep\n\n        if self.ocr_enabled:\n            try:\n                from pdf2image import convert_from_path\n                import pytesseract\n\n                self.convert_from_path = convert_from_path\n                self.pytesseract = pytesseract\n                self.pytesseract.get_tesseract_version()\n            except ImportError:\n                raise ImportError(\n                    \"pdf2image and pytesseract are required to use the pdf transform with ocr. Please install pdf2image and pytesseract with the following command: pip install pdf2image pytesseract\"\n                )\n            except EnvironmentError:\n                raise EnvironmentError(\n                    \"The tesseract engine is required to use the pdf transform with ocr. Please see https://tesseract-ocr.github.io/tessdoc/Installation.html for installation instructions.\"\n                )\n        else:\n            try:\n                from langchain.document_loaders import PDFPlumberLoader\n\n                self.PDFPlumberLoader = PDFPlumberLoader\n            except ImportError:\n                raise ImportError(\n                    \"pdfplumber is required to use the pdf transform. Please install pdfplumber with the following command: pip install pdfplumber\"\n                )\n\n    @staticmethod\n    def name() -> str:\n        return TransformType.PDF\n\n    def get_page_texts(self, row: Dict[str, Any]) -> List[str]:\n        \"\"\"This function gets the text from each page of a PDF file.\n        If OCR is enabled, it uses the pdf2image library to convert the PDF into images and then uses\n        pytesseract to convert the images into text. Otherwise, it uses pdfplumber to extract the text.\n\n        Args:\n            row (Dict[str, Any]): The row of data to be transformed.\n\n        Returns:\n            List[str]: A list of strings containing the text from each page of the PDF.\n        \"\"\"\n        if self.ocr_enabled:\n            pages = self.convert_from_path(row[self.file_path_column])\n            return [self.pytesseract.image_to_string(page) for page in pages]\n        else:\n            loader = self.PDFPlumberLoader(row[self.file_path_column])\n            return [page.page_content for page in loader.load()]\n\n    async def _apply(self, row: Dict[str, Any]) -> Dict[str, Any]:\n        \"\"\"This function transforms a PDF file into a string of text.\n        The text is formatted according to the page_format and\n        page_sep parameters and returned as a string.\n\n        Args:\n            row (Dict[str, Any]): The row of data to be transformed.\n\n        Returns:\n            Dict[str, Any]: The dict of output columns.\n        \"\"\"\n        texts = []\n        for idx, text in enumerate(self.get_page_texts(row)):\n            texts.append(self.page_format.format(page_num=idx + 1, page_content=text))\n        output = self.page_sep.join(texts)\n        transformed_row = {\n            self.output_columns[\"content_column\"]: output,\n            self.output_columns[\"metadata_column\"]: {\"num_pages\": len(texts)},\n        }\n        return self._return_output_row(transformed_row)\n\n    def params(self):\n        return {\n            \"file_path_column\": self.file_path_column,\n            \"ocr_enabled\": self.ocr_enabled,\n            \"page_header\": self.page_format,\n            \"page_sep\": self.page_sep,\n            \"output_columns\": self.output_columns,\n        }\n", "path": "src/autolabel/transforms/pdf.py"}, {"content": "from abc import ABC, abstractmethod\nfrom typing import Dict, Any\nfrom autolabel.cache import BaseCache\nfrom autolabel.schema import TransformCacheEntry\n\n\nclass BaseTransform(ABC):\n    \"\"\"Base class for all transforms.\"\"\"\n\n    TTL_MS = 60 * 60 * 24 * 7 * 1000  # 1 week\n\n    def __init__(self, cache: BaseCache, output_columns: Dict[str, Any]) -> None:\n        \"\"\"\n        Initialize a transform.\n        Args:\n            cache: A cache object to use for caching the results of this transform.\n            output_columns: A dictionary of output columns. The keys are the names of the output columns as expected by the transform. The values are the column names they should be mapped to in the dataset.\n        \"\"\"\n        super().__init__()\n        self._output_columns = output_columns\n        self.cache = cache\n\n    @staticmethod\n    @abstractmethod\n    def name() -> str:\n        \"\"\"\n        Returns the name of the transform.\n        \"\"\"\n        pass\n\n    @property\n    def output_columns(self) -> Dict[str, Any]:\n        \"\"\"\n        Returns a dictionary of output columns. The keys are the names of the output columns\n        as expected by the transform. The values are the column names they should be mapped to in\n        the dataset.\n        \"\"\"\n        return {k: self._output_columns.get(k, None) for k in self.COLUMN_NAMES}\n\n    @abstractmethod\n    async def _apply(self, row: Dict[str, Any]) -> Dict[str, Any]:\n        \"\"\"\n        Applies the transform to the given row.\n        Args:\n            row: A dictionary representing a row in the dataset. The keys are the column names and the values are the column values.\n        Returns:\n            A dictionary representing the transformed row. The keys are the column names and the values are the column values.\n        \"\"\"\n        pass\n\n    @abstractmethod\n    def params(self) -> Dict[str, Any]:\n        \"\"\"\n        Returns a dictionary of parameters that can be used to uniquely identify this transform.\n        Returns:\n            A dictionary of parameters that can be used to uniquely identify this transform.\n        \"\"\"\n        return {}\n\n    async def apply(self, row: Dict[str, Any]) -> Dict[str, Any]:\n        if self.cache is not None:\n            cache_entry = TransformCacheEntry(\n                transform_name=self.name(),\n                transform_params=self.params(),\n                input=row,\n                ttl_ms=self.TTL_MS,\n            )\n            output = self.cache.lookup(cache_entry)\n\n            if output is not None:\n                # Cache hit\n                return output\n\n        output = await self._apply(row)\n\n        if self.cache is not None:\n            cache_entry.output = output\n            self.cache.update(cache_entry)\n\n        return output\n\n    def _return_output_row(self, row: Dict[str, Any]) -> Dict[str, Any]:\n        \"\"\"\n        Returns the output row with the correct column names.\n        Args:\n            row: The output row.\n        Returns:\n            The output row with the correct column names.\n        \"\"\"\n        # remove null key\n        row.pop(None, None)\n        return row\n", "path": "src/autolabel/transforms/base.py"}]}