Source code for langchain.llms.huggingface_text_gen_inference
import logging
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional
from langchain.callbacks.manager import (
    AsyncCallbackManagerForLLMRun,
    CallbackManagerForLLMRun,
)
from langchain.llms.base import LLM
from langchain.pydantic_v1 import Extra, Field, root_validator
from langchain.schema.output import GenerationChunk
from langchain.utils import get_pydantic_field_names
logger = logging.getLogger(__name__)
[docs]class HuggingFaceTextGenInference(LLM):
    """
    HuggingFace text generation API.
    To use, you should have the `text-generation` python package installed and
    a text-generation server running.
    Example:
        .. code-block:: python
            # Basic Example (no streaming)
            llm = HuggingFaceTextGenInference(
                inference_server_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,
            )
            print(llm("What is Deep Learning?"))
            # Streaming response example
            from langchain.callbacks import streaming_stdout
            callbacks = [streaming_stdout.StreamingStdOutCallbackHandler()]
            llm = HuggingFaceTextGenInference(
                inference_server_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,
                callbacks=callbacks,
                streaming=True
            )