huggingface_endpoint.py — langchain Source File

Source Code

from __future__ import annotations

import inspect
import logging
import os
from collections.abc import AsyncIterator, Iterator, Mapping
from typing import Any

from langchain_core.callbacks import (
    AsyncCallbackManagerForLLMRun,
    CallbackManagerForLLMRun,
)
from langchain_core.language_models.llms import LLM
from langchain_core.outputs import GenerationChunk
from langchain_core.utils import from_env, get_pydantic_field_names
from pydantic import ConfigDict, Field, model_validator
from typing_extensions import Self

logger = logging.getLogger(__name__)

VALID_TASKS = (
    "text2text-generation",
    "text-generation",
    "summarization",
    "conversational",
)


class HuggingFaceEndpoint(LLM):
    """Hugging Face Endpoint. This works with any model that supports text generation (i.e. text completion) task.

    To use this class, you should have installed the `huggingface_hub` package, and
    the environment variable `HUGGINGFACEHUB_API_TOKEN` set with your API token,
    or given as a named parameter to the constructor.

    Example:
        ```python
        # Basic Example (no streaming)
        model = HuggingFaceEndpoint(
            endpoint_url="http://localhost:8010/",
            max_new_tokens=512,
            top_k=10,
            top_p=0.95,
            typical_p=0.95,
            temperature=0.01,
            repetition_penalty=1.03,
            huggingfacehub_api_token="my-api-key",
        )
        print(model.invoke("What is Deep Learning?"))

        # Streaming response example
        from langchain_core.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

        callbacks = [StreamingStdOutCallbackHandler()]
        model = HuggingFaceEndpoint(
            endpoint_url="http://localhost:8010/",
            max_new_tokens=512,
            top_k=10,
            top_p=0.95,
            typical_p=0.95,
// ... (396 more lines)