Source code for langchain.llms.rwkv
"""RWKV models.
Based on https://github.com/saharNooby/rwkv.cpp/blob/master/rwkv/chat_with_bot.py
         https://github.com/BlinkDL/ChatRWKV/blob/main/v2/chat.py
"""
from typing import Any, Dict, List, Mapping, Optional, Set
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.llms.base import LLM
from langchain.llms.utils import enforce_stop_tokens
from langchain.pydantic_v1 import BaseModel, Extra, root_validator
[docs]class RWKV(LLM, BaseModel):
    """RWKV language models.
    To use, you should have the ``rwkv`` python package installed, the
    pre-trained model file, and the model's config information.
    Example:
        .. code-block:: python
            from langchain.llms import RWKV
            model = RWKV(model="./models/rwkv-3b-fp16.bin", strategy="cpu fp32")
            # Simplest invocation
            response = model("Once upon a time, ")
    """
    model: str
    """Path to the pre-trained RWKV model file."""
    tokens_path: str
    """Path to the RWKV tokens file."""
    strategy: str = "cpu fp32"
    """Token context window."""
    rwkv_verbose: bool = True
    """Print debug information."""
    temperature: float = 1.0
    """The temperature to use for sampling."""
    top_p: float = 0.5
    """The top-p value to use for sampling."""
    penalty_alpha_frequency: float = 0.4
    """Positive values penalize new tokens based on their existing frequency