langextract-vllm 0.1.0__py3-none-any.whl

langextract_vllm/__init__.py

@@ -0,0 +1,6 @@
+"""LangExtract provider plugin for VLLM."""
+
+from langextract_vllm.provider import VLLMLanguageModel
+
+__all__ = ['VLLMLanguageModel']
+__version__ = "0.1.0"

langextract_vllm/provider.py

@@ -0,0 +1,185 @@
+"""Provider implementation for vLLM (Direct Library Integration Only)."""
+
+from typing import List, Generator, Dict, Any
+
+import langextract as lx
+from langextract.core import base_model, exceptions, types
+
+VLLM_PATTERNS = (
+    r"^vllm:",
+)
+
+
+@lx.providers.registry.register(*VLLM_PATTERNS, priority=20)
+class VLLMLanguageModel(base_model.BaseLanguageModel):
+    """LangExtract provider for direct vLLM library integration.
+
+    This provider handles model IDs matching: ['^vllm:']
+    Uses vLLM library directly for optimal performance.
+    """
+
+    def __init__(
+        self,
+        model_id: str,
+        max_workers: int = 1,
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        max_tokens: int = 1024,
+        gpu_memory_utilization: float = 0.7,  # 进一步降低GPU内存利用率
+        max_model_len: int = 1024,  # 进一步减少最大序列长度
+        **kwargs,
+    ):
+        """Initialize the direct vLLM provider.
+
+        Args:
+            model_id: Model identifier in format "vllm:model_name_or_path"
+            max_workers: Maximum parallel workers (affects batch size)
+            temperature: Sampling temperature
+            top_p: Top-p sampling parameter
+            max_tokens: Maximum tokens to generate
+            gpu_memory_utilization: GPU memory utilization ratio (0.0-1.0)
+            max_model_len: Maximum sequence length (reduces KV cache memory)
+            **kwargs: Additional vLLM initialization parameters
+        """
+        try:
+            from vllm import LLM, SamplingParams
+        except ImportError:
+            raise exceptions.InferenceConfigError(
+                "vLLM library not installed. Please install with: pip install vllm"
+            )
+
+        super().__init__()
+
+        self.model_id = model_id
+        self.max_workers = max_workers
+        self.gpu_memory_utilization = gpu_memory_utilization
+        self.max_model_len = max_model_len
+
+        # Extract actual model name from identifier
+        if model_id.startswith("vllm:"):
+            actual_model_id = model_id.split(":", 1)[1]
+        else:
+            actual_model_id = model_id
+
+        # Separate vLLM engine kwargs from sampling kwargs
+        engine_kwargs = {
+            k: v for k, v in kwargs.items() 
+            if k not in ['temperature', 'top_p', 'max_tokens', 'gpu_memory_utilization', 'max_model_len']
+        }
+
+        # Initialize vLLM engine with memory control
+        # Disable torch.compile to avoid PY_SSIZE_T_CLEAN error
+        engine_kwargs.setdefault('enforce_eager', True)
+        engine_kwargs.setdefault('disable_custom_all_reduce', True)
+        
+        self._vllm_engine = LLM(
+            model=actual_model_id,
+            gpu_memory_utilization=gpu_memory_utilization,
+            max_model_len=max_model_len,
+            **engine_kwargs
+        )
+
+        # Configure sampling parameters
+        self._sampling_params = SamplingParams(
+            temperature=temperature,
+            top_p=top_p,
+            max_tokens=max_tokens,
+        )
+
+    def _process_batch(self, prompts: List[str]) -> List[types.ScoredOutput]:
+        """Process batch using direct vLLM engine."""
+        try:
+            outputs = self._vllm_engine.generate(
+                prompts,
+                self._sampling_params,
+                use_tqdm=False
+            )
+
+            results = []
+            for output in outputs:
+                result = output.outputs[0].text
+                results.append(types.ScoredOutput(score=1.0, output=result))
+
+            return results
+
+        except Exception as e:
+            raise exceptions.InferenceRuntimeError(
+                f"vLLM engine error: {str(e)}", original=e
+            ) from e
+
+    def infer(self, batch_prompts: List[str], **kwargs) -> Generator[List[types.ScoredOutput], None, None]:
+        """Run inference using direct vLLM integration.
+
+        Args:
+            batch_prompts: List of prompts to process.
+            **kwargs: Additional sampling parameters (temperature, top_p, max_tokens).
+
+        Yields:
+            Lists of ScoredOutput objects, one per batch.
+        """
+        # Update sampling parameters if provided
+        sampling_params = self._sampling_params
+        if kwargs:
+            sampling_kwargs = {
+                k: kwargs[k] for k in ['temperature', 'top_p', 'max_tokens'] 
+                if k in kwargs
+            }
+            if sampling_kwargs:
+                # Create new sampling params with updated values
+                current_params = {
+                    'temperature': sampling_params.temperature,
+                    'top_p': sampling_params.top_p,
+                    'max_tokens': sampling_params.max_tokens,
+                }
+                current_params.update(sampling_kwargs)
+                
+                from vllm import SamplingParams
+                sampling_params = SamplingParams(**current_params)
+
+        # Store original sampling params and update for this inference
+        original_params = self._sampling_params
+        self._sampling_params = sampling_params
+
+        try:
+            # Process in batches for better memory management
+            # Adjust batch size based on max_workers and available resources
+            batch_size = max(1, self.max_workers * 4)
+
+            for i in range(0, len(batch_prompts), batch_size):
+                batch_chunk = batch_prompts[i:i + batch_size]
+                chunk_results = self._process_batch(batch_chunk)
+                yield chunk_results
+
+        finally:
+            # Restore original sampling parameters
+            self._sampling_params = original_params
+
+    @classmethod
+    def get_supported_models(cls) -> List[Dict[str, Any]]:
+        """Get list of supported models and their capabilities."""
+        return [
+            {
+                "id": "vllm:*",
+                "name": "vLLM Direct",
+                "description": "Direct vLLM library integration with PagedAttention support",
+                "capabilities": ["chat", "completion", "batch_processing"],
+                "max_tokens": 131072,  # 128k context window (model dependent)
+                "supports_parallel": True,
+                "supports_batching": True,
+            }
+        ]
+
+    def get_model_info(self) -> Dict[str, Any]:
+        """Get information about the current model configuration."""
+        return {
+            "model_id": self.model_id,
+            "max_workers": self.max_workers,
+            "gpu_memory_utilization": self.gpu_memory_utilization,
+            "max_model_len": self.max_model_len,
+            "sampling_params": {
+                "temperature": self._sampling_params.temperature,
+                "top_p": self._sampling_params.top_p,
+                "max_tokens": self._sampling_params.max_tokens,
+            },
+            "provider_type": "vllm_direct",
+        }

langextract_vllm-0.1.0.dist-info/METADATA

@@ -0,0 +1,46 @@
+Metadata-Version: 2.4
+Name: langextract-vllm
+Version: 0.1.0
+Summary: LangExtract provider plugin for VLLM - 高性能大语言模型推理插件
+Project-URL: Homepage, https://github.com/langextract/langextract-vllm
+Project-URL: Repository, https://github.com/langextract/langextract-vllm
+Project-URL: Documentation, https://langextract-vllm.readthedocs.io
+Project-URL: Bug Tracker, https://github.com/langextract/langextract-vllm/issues
+Project-URL: Changelog, https://github.com/langextract/langextract-vllm/blob/main/CHANGELOG.md
+Author-email: LangExtract Team <contact@langextract.dev>
+Maintainer-email: LangExtract Team <contact@langextract.dev>
+License: MIT
+License-File: LICENSE
+Keywords: extraction,gpu,inference,langextract,llm,nlp,pytorch,vllm
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Topic :: Text Processing :: Linguistic
+Requires-Python: >=3.10
+Requires-Dist: langextract>=1.0.0
+Requires-Dist: torch>=2.0.0
+Requires-Dist: transformers>=4.30.0
+Requires-Dist: vllm>=0.2.0
+Provides-Extra: dev
+Requires-Dist: black>=23.0.0; extra == 'dev'
+Requires-Dist: flake8>=6.0.0; extra == 'dev'
+Requires-Dist: isort>=5.12.0; extra == 'dev'
+Requires-Dist: mypy>=1.0.0; extra == 'dev'
+Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
+Requires-Dist: pytest>=7.0.0; extra == 'dev'
+Provides-Extra: docs
+Requires-Dist: myst-parser>=1.0.0; extra == 'docs'
+Requires-Dist: sphinx-rtd-theme>=1.2.0; extra == 'docs'
+Requires-Dist: sphinx>=5.0.0; extra == 'docs'
+Provides-Extra: test
+Requires-Dist: pytest-cov>=4.0.0; extra == 'test'
+Requires-Dist: pytest-mock>=3.10.0; extra == 'test'
+Requires-Dist: pytest>=7.0.0; extra == 'test'

langextract_vllm-0.1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+langextract_vllm/__init__.py,sha256=z-S-_mwU-g-J_psIkmrCZLmwG76l6KhLNpS071U8KrM,156
+langextract_vllm/provider.py,sha256=Hqs2vzZYadwmIG-VOnAvodGIthjtOzUV9TpgzULvnPs,6857
+langextract_vllm-0.1.0.dist-info/METADATA,sha256=Vuf-Rwgn73y4XK27n_As2NBP9zwfnexnyznMeTmdzSQ,2169
+langextract_vllm-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+langextract_vllm-0.1.0.dist-info/entry_points.txt,sha256=3M2Y4ac8Sd2wJaga8rXJT5Lnk8kUmwGmhnsiSQvhc7o,145
+langextract_vllm-0.1.0.dist-info/licenses/LICENSE,sha256=20V09KLo5sLIeUpiJCTpgvsKqX0uPQrwhaGK1GYb_LM,1069
+langextract_vllm-0.1.0.dist-info/RECORD,,

langextract_vllm-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

langextract_vllm-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,5 @@
+[console_scripts]
+langextract-vllm-test = langextract_vllm.test:main
+
+[langextract.providers]
+vllm = langextract_vllm.provider:VLLMLanguageModel

langextract_vllm-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Fajar Garnadi
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.