onecrawler 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (50) hide show
  1. onecrawler-0.1.0/LICENSE +21 -0
  2. onecrawler-0.1.0/PKG-INFO +22 -0
  3. onecrawler-0.1.0/README.md +309 -0
  4. onecrawler-0.1.0/onecrawler/__init__.py +7 -0
  5. onecrawler-0.1.0/onecrawler/browser.py +93 -0
  6. onecrawler-0.1.0/onecrawler/crawler/__init__.py +0 -0
  7. onecrawler-0.1.0/onecrawler/crawler/base.py +54 -0
  8. onecrawler-0.1.0/onecrawler/crawler/link/__init__.py +0 -0
  9. onecrawler-0.1.0/onecrawler/crawler/link/classifier.py +116 -0
  10. onecrawler-0.1.0/onecrawler/crawler/link/deep.py +235 -0
  11. onecrawler-0.1.0/onecrawler/crawler/link/engine.py +80 -0
  12. onecrawler-0.1.0/onecrawler/crawler/link/helper.py +124 -0
  13. onecrawler-0.1.0/onecrawler/crawler/link/shallow.py +90 -0
  14. onecrawler-0.1.0/onecrawler/crawler/scraper/__init__.py +0 -0
  15. onecrawler-0.1.0/onecrawler/crawler/scraper/engine.py +90 -0
  16. onecrawler-0.1.0/onecrawler/crawler/scraper/genai/__init__.py +0 -0
  17. onecrawler-0.1.0/onecrawler/crawler/scraper/genai/executor.py +39 -0
  18. onecrawler-0.1.0/onecrawler/crawler/scraper/genai/graph.py +61 -0
  19. onecrawler-0.1.0/onecrawler/crawler/scraper/genai/model.py +44 -0
  20. onecrawler-0.1.0/onecrawler/crawler/scraper/genai/prompt.py +22 -0
  21. onecrawler-0.1.0/onecrawler/crawler/scraper/genai/state.py +9 -0
  22. onecrawler-0.1.0/onecrawler/crawler/scraper/heuristic/script.py +50 -0
  23. onecrawler-0.1.0/onecrawler/map/__init__.py +0 -0
  24. onecrawler-0.1.0/onecrawler/map/helper.py +63 -0
  25. onecrawler-0.1.0/onecrawler/map/sitemap.py +563 -0
  26. onecrawler-0.1.0/onecrawler/proxy/__init__.py +0 -0
  27. onecrawler-0.1.0/onecrawler/proxy/pool.py +28 -0
  28. onecrawler-0.1.0/onecrawler/settings/__init__.py +4 -0
  29. onecrawler-0.1.0/onecrawler/settings/browser.py +73 -0
  30. onecrawler-0.1.0/onecrawler/settings/crawler.py +83 -0
  31. onecrawler-0.1.0/onecrawler/settings/genai.py +17 -0
  32. onecrawler-0.1.0/onecrawler/settings/proxy.py +43 -0
  33. onecrawler-0.1.0/onecrawler/settings/simulation.py +25 -0
  34. onecrawler-0.1.0/onecrawler/utils/__init__.py +0 -0
  35. onecrawler-0.1.0/onecrawler/utils/decorator.py +24 -0
  36. onecrawler-0.1.0/onecrawler/utils/stats.py +0 -0
  37. onecrawler-0.1.0/onecrawler/utils/writter.py +6 -0
  38. onecrawler-0.1.0/onecrawler.egg-info/PKG-INFO +22 -0
  39. onecrawler-0.1.0/onecrawler.egg-info/SOURCES.txt +48 -0
  40. onecrawler-0.1.0/onecrawler.egg-info/dependency_links.txt +1 -0
  41. onecrawler-0.1.0/onecrawler.egg-info/requires.txt +16 -0
  42. onecrawler-0.1.0/onecrawler.egg-info/top_level.txt +3 -0
  43. onecrawler-0.1.0/pyproject.toml +44 -0
  44. onecrawler-0.1.0/setup.cfg +4 -0
  45. onecrawler-0.1.0/tests/test_deep_crawler.py +121 -0
  46. onecrawler-0.1.0/tests/test_link_helpers.py +48 -0
  47. onecrawler-0.1.0/tests/test_map_helper.py +51 -0
  48. onecrawler-0.1.0/tests/test_settings.py +76 -0
  49. onecrawler-0.1.0/tests/test_sitemap.py +188 -0
  50. onecrawler-0.1.0/tests/test_utils_and_base.py +84 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Sayed Shaun
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,22 @@
1
+ Metadata-Version: 2.4
2
+ Name: onecrawler
3
+ Version: 0.1.0
4
+ Summary: A feature-rich web crawler package built with Python.
5
+ Author: Sayed Shaun
6
+ Requires-Python: >=3.10
7
+ License-File: LICENSE
8
+ Requires-Dist: aiohttp
9
+ Requires-Dist: curl_cffi
10
+ Requires-Dist: tqdm
11
+ Requires-Dist: playwright
12
+ Requires-Dist: pydantic
13
+ Requires-Dist: trafilatura
14
+ Requires-Dist: lxml
15
+ Provides-Extra: genai
16
+ Requires-Dist: langchain; extra == "genai"
17
+ Requires-Dist: langgraph; extra == "genai"
18
+ Provides-Extra: dev
19
+ Requires-Dist: black; extra == "dev"
20
+ Requires-Dist: isort; extra == "dev"
21
+ Requires-Dist: pre-commit; extra == "dev"
22
+ Dynamic: license-file
@@ -0,0 +1,309 @@
1
+ <div align="center">
2
+
3
+ <img src="https://raw.githubusercontent.com/sayedshaun/onecrawler/refs/heads/main/docs/static/onecrawl_logo.png" alt="Onecrawler" width="200"/>
4
+
5
+ # Onecrawler
6
+
7
+ **An async Python toolkit for sitemap discovery, browser crawling, and structured content extraction.**
8
+
9
+ [![CI](https://github.com/sayedshaun/onecrawler/actions/workflows/ci.yml/badge.svg)](https://github.com/sayedshaun/onecrawler/actions/workflows/ci.yml)
10
+ [![Python](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://www.python.org/downloads/)
11
+ [![PyPI](https://img.shields.io/badge/pypi-onecrawler-2ea44f.svg)](https://pypi.org/project/onecrawler/)
12
+ [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
13
+ [![Imports: isort](https://img.shields.io/badge/imports-isort-1674b1.svg)](https://pycqa.github.io/isort/)
14
+ [![License: MIT](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE)
15
+
16
+ [Installation](#installation) · [Quick Start](#quick-start) · [Documentation](#documentation) · [Development](#development)
17
+
18
+ </div>
19
+
20
+ ---
21
+
22
+ ## Overview
23
+
24
+ Onecrawler helps you build maintainable crawling and extraction workflows without
25
+ turning every project into a custom scraping script. It gives you a shared
26
+ settingsuration model, async execution, sitemap discovery, browser-backed link
27
+ extraction, heuristic content extraction, and optional GenAI extraction for typed
28
+ outputs.
29
+
30
+ The recommended workflow is:
31
+
32
+ 1. Use sitemaps first whenever possible.
33
+ 2. Fall back to browser link extraction when sitemap coverage is missing or dynamic.
34
+ 3. Scrape the final URL list with heuristic extraction by default.
35
+ 4. Use GenAI extraction when you need structured output in a Pydantic schema.
36
+
37
+ ```python
38
+ sitemap = UniversalSiteMap(settings)
39
+ urls = await sitemap.run("https://example.com")
40
+
41
+ async with ScraperEngine(settings) as scraper:
42
+ records = await scraper.run(urls)
43
+ ```
44
+
45
+ ---
46
+
47
+ ## Features
48
+
49
+ | Capability | Details |
50
+ | --- | --- |
51
+ | **Sitemap discovery** | Resolves `robots.txt`, common sitemap paths, nested indexes, `.xml.gz`, feeds, and HTML fallback |
52
+ | **Browser link extraction** | Shallow and deep Playwright-backed discovery for JavaScript-rendered or sitemap-poor sites |
53
+ | **URL filtering** | Wildcard path filters with `include_link_patterns` |
54
+ | **Async performance** | Tunable concurrency, retries, timeouts, and crawl limits |
55
+ | **Content extraction** | Heuristic extraction with `trafilatura` for fast article-like content extraction |
56
+ | **GenAI extraction** | Optional model-assisted extraction for strongly typed Pydantic outputs |
57
+ | **Output formats** | `markdown`, `json`, `csv`, `html`, `python`, `txt`, `xml`, `xmltei` |
58
+ | **Proxy support** | Single proxy or rotating proxy pools for browser and sitemap workflows |
59
+ | **Browser controls** | Viewport, user agent, locale, timezone, storage state, and runtime settings |
60
+
61
+ ---
62
+
63
+ ## When To Use What
64
+
65
+ | Need | Use | Why |
66
+ | --- | --- | --- |
67
+ | Fast URL discovery from a public site | `UniversalSiteMap` | It is usually the simplest, fastest, and least expensive way to collect URLs |
68
+ | Links from one listing page | Shallow `LinkExtractionEngine` | It reads direct same-site links from the page |
69
+ | Recursive discovery through navigation | Deep `LinkExtractionEngine` | It follows internal links until your settingsured limit |
70
+ | Bulk article or page text extraction | Heuristic `ScraperEngine` | It is deterministic and avoids model cost |
71
+ | Typed fields or semantic normalization | GenAI extraction | It can produce schema-shaped output for downstream systems |
72
+
73
+ ---
74
+
75
+ ## Installation
76
+
77
+ ```bash
78
+ pip install onecrawler
79
+ ```
80
+
81
+ Install Playwright browser binaries when you use browser-backed crawling or scraping:
82
+
83
+ ```bash
84
+ python -m playwright install chromium
85
+ ```
86
+
87
+ Install optional GenAI dependencies when you use model-assisted extraction:
88
+
89
+ ```bash
90
+ pip install "onecrawler[genai]"
91
+ ```
92
+
93
+ For local development:
94
+
95
+ ```bash
96
+ git clone https://github.com/sayedshaun/onecrawler.git
97
+ cd onecrawler
98
+ python -m pip install -e ".[dev]"
99
+ python -m playwright install chromium
100
+ ```
101
+
102
+ ---
103
+
104
+ ## Quick Start
105
+
106
+ This example uses the production-friendly path: discover URLs from the sitemap, then
107
+ scrape them.
108
+
109
+ ```python
110
+ import json
111
+ import asyncio
112
+ from onecrawler import CrawlerSettings, ScraperEngine, UniversalSiteMap
113
+
114
+
115
+ async def main():
116
+ settings = CrawlerSettings(
117
+ link_extraction_limit=100,
118
+ include_link_patterns=["/articles/*"],
119
+ scraping_strategy="heuristic",
120
+ scraping_output_format="json",
121
+ concurrency=8,
122
+ request_timeout=15,
123
+ max_retries=3,
124
+ )
125
+
126
+ sitemap = UniversalSiteMap(settings)
127
+ urls = await sitemap.run("https://example.com")
128
+
129
+ async with ScraperEngine(settings) as scraper:
130
+ records = await scraper.run(urls)
131
+
132
+ with open("articles.json", "w", encoding="utf-8") as f:
133
+ json.dump(records, f, indent=2, ensure_ascii=False)
134
+
135
+
136
+ if __name__ == "__main__":
137
+ asyncio.run(main())
138
+ ```
139
+
140
+ ### Browser Link Extraction
141
+
142
+ Use browser extraction when sitemaps are incomplete, unavailable, or unable to expose
143
+ JavaScript-rendered links.
144
+
145
+ ```python
146
+ import asyncio
147
+ from onecrawler import CrawlerSettings, LinkExtractionEngine
148
+
149
+
150
+ async def main():
151
+ settings = CrawlerSettings(
152
+ link_extraction_strategy="deep",
153
+ link_extraction_limit=250,
154
+ include_link_patterns=["/news/*"],
155
+ concurrency=5,
156
+ )
157
+
158
+ async with LinkExtractionEngine(settings) as engine:
159
+ links = await engine.run("https://example.com/news")
160
+
161
+ print(f"Collected {len(links)} links")
162
+
163
+
164
+ if __name__ == "__main__":
165
+ asyncio.run(main())
166
+ ```
167
+
168
+ ### GenAI Extraction With A Schema
169
+
170
+ Use GenAI extraction when you need a strongly typed response shape instead of plain
171
+ content. This requires installing the GenAI dependencies:
172
+
173
+ ```bash
174
+ pip install "onecrawler[genai]"
175
+ ```
176
+
177
+ ```python
178
+ import asyncio
179
+ from typing import Optional
180
+ from pydantic import BaseModel
181
+ from onecrawler import CrawlerSettings, GenerativeAISettings, ScraperEngine
182
+
183
+
184
+ class ArticleSummary(BaseModel):
185
+ title: str
186
+ author: Optional[str] = None
187
+ published_at: Optional[str] = None
188
+ summary: str
189
+ topics: list[str]
190
+
191
+
192
+ async def main():
193
+ settings = CrawlerSettings(
194
+ scraping_strategy="genai", # Required for GenAI extraction
195
+ scraping_output_format="json", # GenAI only supports JSON
196
+ genai=GenerativeAISettings(
197
+ provider="openai", # Options: "openai", "google", "ollama"
198
+ model_name="gpt-4o-mini",
199
+ api_key="YOUR_API_KEY", # Required for OpenAI/Google, optional for Ollama
200
+ output_schema=ArticleSummary, # Pydantic model for structured output
201
+ # Optional: base_url for custom endpoints (e.g., Ollama)
202
+ # base_url="https://your-ollama-instance.com/",
203
+ ),
204
+ concurrency=2, # Lower concurrency recommended for GenAI
205
+ request_timeout=30, # Increase timeout for model responses
206
+ )
207
+
208
+ async with ScraperEngine(settings) as scraper:
209
+ result = await scraper.run("https://example.com/articles/story")
210
+
211
+ # Convert Pydantic model to dict for JSON serialization
212
+ print(result.model_dump() if hasattr(result, 'model_dump') else result)
213
+
214
+
215
+ if __name__ == "__main__":
216
+ asyncio.run(main())
217
+ ```
218
+
219
+ #### Supported Providers
220
+
221
+ - **OpenAI**: Requires `api_key`, supports GPT models
222
+ - **Google**: Requires `api_key`, supports Gemini models
223
+ - **Ollama**: No API key needed, requires `base_url`, supports local models
224
+
225
+ #### Ollama Example
226
+
227
+ ```python
228
+ settings = CrawlerSettings(
229
+ scraping_strategy="genai",
230
+ genai=GenerativeAISettings(
231
+ provider="ollama",
232
+ model_name="llama3:8b",
233
+ base_url="http://localhost:11434/", # Your Ollama instance
234
+ output_schema=ArticleSummary,
235
+ ),
236
+ )
237
+ ```
238
+
239
+ ### Proxy Support
240
+
241
+ Attach one proxy or a rotating proxy pool directly to `CrawlerSettings`.
242
+
243
+ ```python
244
+ from onecrawler import CrawlerSettings, ProxySettings
245
+
246
+
247
+ settings = CrawlerSettings(
248
+ proxies=[
249
+ ProxySettings(server="http://proxy-1.example:8080"),
250
+ ProxySettings(
251
+ server="http://proxy-2.example:8080",
252
+ username="user",
253
+ password="pass",
254
+ ),
255
+ ],
256
+ proxy_rotation="round_robin",
257
+ )
258
+ ```
259
+
260
+ Use `proxy=ProxySettings(...)` for one proxy, or `proxies=[...]` for rotation.
261
+
262
+ ---
263
+
264
+ ## Documentation
265
+
266
+ The README is the project overview. The full documentation in [`docs/`](docs/index.md)
267
+ contains production guidance, caveats, performance notes, and copy-paste examples.
268
+
269
+ | Topic | Guide |
270
+ | --- | --- |
271
+ | Install the package | [Installation](docs/installation.md) |
272
+ | Run your first crawl | [Quick start](docs/quick-start.md) |
273
+ | Tune crawler settings | [settings](docs/settings.md) |
274
+ | Discover URLs from sitemaps | [Sitemap discovery](docs/sitemap-discovery.md) |
275
+ | Extract and filter links | [Link extraction](docs/link-extraction.md) |
276
+ | Scrape page content | [Scraping](docs/scraping.md) |
277
+ | Public classes and exports | [API reference](docs/api-reference.md) |
278
+ | Common fixes | [Troubleshooting](docs/troubleshooting.md) |
279
+ | Contribute locally | [Contributing](docs/contributing.md) |
280
+ | Work on the project | [Development](docs/development.md) |
281
+
282
+ See [Contributing](docs/contributing.md) for how to improve the docs.
283
+
284
+ ---
285
+
286
+ ## Production Tips
287
+
288
+ - Prefer `UniversalSiteMap` before browser crawling.
289
+ - Always set `link_extraction_limit` for broad jobs.
290
+ - Use `include_link_patterns` to keep discovery focused.
291
+ - Start with moderate `concurrency`, then increase gradually.
292
+ - Use heuristic scraping for bulk content extraction.
293
+ - Use GenAI extraction for schema-shaped output, summaries, classification, or field
294
+ normalization.
295
+ - Split discovery and scraping into separate steps for easier retries.
296
+
297
+ ---
298
+
299
+ ## License
300
+
301
+ Released under the [MIT License](LICENSE). See `LICENSE` for full terms.
302
+
303
+ ---
304
+
305
+ <div align="center">
306
+
307
+ Built by [sayedshaun](https://github.com/sayedshaun)
308
+
309
+ </div>
@@ -0,0 +1,7 @@
1
+ from .crawler.link.classifier import LinkClassifierPipeline
2
+ from .crawler.link.engine import LinkExtractionEngine
3
+ from .crawler.scraper.engine import ScraperEngine
4
+ from .map.sitemap import SiteMap, SitemapStats, UniversalSiteMap
5
+ from .settings import *
6
+
7
+ __version__ = "0.1.0"
@@ -0,0 +1,93 @@
1
+ from contextlib import suppress
2
+
3
+ from playwright.async_api import async_playwright
4
+
5
+ from .settings.browser import BrowserSettings
6
+
7
+
8
+ class GoogleChrome:
9
+ def __init__(self, settings: BrowserSettings):
10
+ self.settings = settings
11
+ self.playwright = None
12
+ self.browser = None
13
+ self.context = None
14
+ self._started = False
15
+ self._closed = False
16
+
17
+ async def start(self):
18
+ if self._started:
19
+ return
20
+
21
+ self.playwright = await async_playwright().start()
22
+
23
+ launch = self.settings.launch
24
+ context = self.settings.context
25
+
26
+ self.browser = await self.playwright.chromium.launch(
27
+ headless=launch.headless,
28
+ slow_mo=launch.slow_mo,
29
+ args=launch.args,
30
+ executable_path=launch.executable_path,
31
+ channel=launch.channel,
32
+ env=launch.env,
33
+ )
34
+
35
+ self.context = await self.browser.new_context(
36
+ viewport=context.viewport,
37
+ screen=context.screen,
38
+ no_viewport=context.no_viewport,
39
+ locale=context.locale,
40
+ timezone_id=context.timezone_id,
41
+ user_agent=context.user_agent,
42
+ java_script_enabled=context.java_script_enabled,
43
+ bypass_csp=context.bypass_csp,
44
+ ignore_https_errors=context.ignore_https_errors,
45
+ extra_http_headers=context.extra_http_headers,
46
+ offline=context.offline,
47
+ geolocation=context.geolocation,
48
+ permissions=context.permissions,
49
+ storage_state=context.storage_state,
50
+ base_url=context.base_url,
51
+ proxy=self.settings.proxy.as_playwright() if self.settings.proxy else None,
52
+ )
53
+
54
+ self._started = True
55
+ self._closed = False
56
+
57
+ async def new_page(self):
58
+ if not self._started:
59
+ await self.start()
60
+
61
+ page = await self.context.new_page()
62
+
63
+ runtime = self.settings.runtime
64
+ page.set_default_timeout(runtime.action_timeout)
65
+ page.set_default_navigation_timeout(runtime.navigation_timeout)
66
+
67
+ return page
68
+
69
+ async def close(self):
70
+ if self._closed:
71
+ return
72
+
73
+ self._closed = True
74
+
75
+ # Close context safely
76
+ if self.context:
77
+ with suppress(Exception):
78
+ await self.context.close()
79
+ self.context = None
80
+
81
+ # Close browser (YOU WERE MISSING THIS)
82
+ if self.browser:
83
+ with suppress(Exception):
84
+ await self.browser.close()
85
+ self.browser = None
86
+
87
+ # Stop playwright
88
+ if self.playwright:
89
+ with suppress(Exception):
90
+ await self.playwright.stop()
91
+ self.playwright = None
92
+
93
+ self._started = False
File without changes
@@ -0,0 +1,54 @@
1
+ import logging
2
+ from abc import ABC, abstractmethod
3
+ from typing import Optional, Type
4
+
5
+
6
+ class BaseEngine(ABC):
7
+ def __init__(self):
8
+ self._closed: bool = True
9
+ self.logger = logging.getLogger(self.__class__.__name__)
10
+
11
+ # ===== Context Manager =====
12
+ async def __aenter__(self):
13
+ self._closed = False
14
+ await self.start()
15
+ self.logger.debug("Engine started")
16
+ return self
17
+
18
+ async def __aexit__(
19
+ self,
20
+ exc_type: Optional[Type[BaseException]],
21
+ exc: Optional[BaseException],
22
+ tb,
23
+ ):
24
+ try:
25
+ await self.close()
26
+ finally:
27
+ self._closed = True
28
+ self.logger.debug("Engine closed")
29
+
30
+ # ===== Lifecycle Hooks =====
31
+ async def start(self):
32
+ """Override to initialize resources."""
33
+ pass
34
+
35
+ async def close(self):
36
+ """Override to cleanup resources."""
37
+ pass
38
+
39
+ # ===== REQUIRED API =====
40
+ @abstractmethod
41
+ async def run(self, *args, **kwargs):
42
+ """Main execution method for the engine."""
43
+ raise NotImplementedError
44
+
45
+ # ===== Safety =====
46
+ def _ensure_open(self):
47
+ if self._closed:
48
+ raise RuntimeError(
49
+ f"{self.__class__.__name__} is closed. Use 'async with'."
50
+ )
51
+
52
+ @property
53
+ def is_closed(self) -> bool:
54
+ return self._closed
File without changes
@@ -0,0 +1,116 @@
1
+ import asyncio
2
+ import warnings
3
+ from functools import lru_cache
4
+ from typing import List
5
+ from urllib.parse import unquote
6
+
7
+ CLASSIFIER_AVAILABLE = True
8
+ try:
9
+ import torch
10
+ from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
11
+ except ImportError:
12
+ CLASSIFIER_AVAILABLE = False
13
+
14
+
15
+ if CLASSIFIER_AVAILABLE:
16
+ repo_id = "SayedShaun/distilbert-link-type-classifier"
17
+
18
+ torch.set_grad_enabled(False)
19
+
20
+ model = AutoModelForSequenceClassification.from_pretrained(repo_id)
21
+ tokenizer = AutoTokenizer.from_pretrained(repo_id)
22
+
23
+ def get_classifier(device: str = "cpu"):
24
+ device_id = 0 if device == "cuda" else -1
25
+ return pipeline(
26
+ "text-classification",
27
+ model=model,
28
+ tokenizer=tokenizer,
29
+ device=device_id,
30
+ )
31
+
32
+ clf = get_classifier("cuda" if torch.cuda.is_available() else "cpu")
33
+
34
+
35
+ def cheap_filter(url: str) -> bool:
36
+ if not url:
37
+ return False
38
+
39
+ bad_patterns = (
40
+ "javascript:",
41
+ "mailto:",
42
+ "tel:",
43
+ "#",
44
+ )
45
+
46
+ return not any(p in url for p in bad_patterns)
47
+
48
+
49
+ @lru_cache(maxsize=10000)
50
+ def _cached_single_prediction(url: str) -> str:
51
+ result = clf(url, truncation=True, max_length=128)[0]
52
+ return result["label"]
53
+
54
+
55
+ class LinkClassifierPipeline:
56
+ def __init__(
57
+ self,
58
+ confidence_threshold: float = 0.8,
59
+ ):
60
+ self.threshold = confidence_threshold
61
+
62
+ if not CLASSIFIER_AVAILABLE:
63
+ warnings.warn(
64
+ "Classifier enabled but transformers/torch not installed. Disabling."
65
+ )
66
+ self.available = False
67
+ else:
68
+ self.available = True
69
+
70
+ async def classify_batch(self, urls: List[str]) -> List[bool]:
71
+ if not self.available:
72
+ return [True] * len(urls)
73
+
74
+ filtered_urls = []
75
+ index_map = []
76
+
77
+ for i, url in enumerate(urls):
78
+ if cheap_filter(url):
79
+ filtered_urls.append(unquote(url))
80
+ index_map.append(i)
81
+
82
+ results = [False] * len(urls)
83
+
84
+ if not filtered_urls:
85
+ return results
86
+
87
+ try:
88
+ predictions = await asyncio.to_thread(
89
+ lambda: clf(filtered_urls, truncation=True, max_length=128)
90
+ )
91
+ except Exception:
92
+ return [True] * len(urls)
93
+
94
+ for idx, pred in zip(index_map, predictions):
95
+ label = pred["label"]
96
+ score = pred["score"]
97
+
98
+ if label == "content" and score >= self.threshold:
99
+ results[idx] = True
100
+ else:
101
+ results[idx] = False
102
+
103
+ return results
104
+
105
+ async def is_valid(self, url: str) -> bool:
106
+ if not self.available:
107
+ return True
108
+
109
+ if not cheap_filter(url):
110
+ return False
111
+
112
+ try:
113
+ label = await asyncio.to_thread(_cached_single_prediction, unquote(url))
114
+ return label != "section"
115
+ except Exception:
116
+ return True