import os from typing import Any, Callable, Dict, Optional, Sequence import requests from tqdm import tqdm from llama_index.legacy.bridge.pydantic import Field, PrivateAttr from llama_index.legacy.callbacks import CallbackManager from llama_index.legacy.constants import ( DEFAULT_CONTEXT_WINDOW, DEFAULT_NUM_OUTPUTS, DEFAULT_TEMPERATURE, ) from llama_index.legacy.core.llms.types import ( ChatMessage, ChatResponse, ChatResponseGen, CompletionResponse, CompletionResponseGen, LLMMetadata, ) from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback from llama_index.legacy.llms.custom import CustomLLM from llama_index.legacy.llms.generic_utils import ( completion_response_to_chat_response, stream_completion_response_to_chat_response, ) from llama_index.legacy.types import BaseOutputParser, PydanticProgramMode from llama_index.legacy.utils import get_cache_dir DEFAULT_LLAMA_CPP_GGML_MODEL = ( "https://huggingface.co/TheBloke/Llama-2-13B-chat-GGML/resolve" "/main/llama-2-13b-chat.ggmlv3.q4_0.bin" ) DEFAULT_LLAMA_CPP_GGUF_MODEL = ( "https://huggingface.co/TheBloke/Llama-2-13B-chat-GGUF/resolve" "/main/llama-2-13b-chat.Q4_0.gguf" ) DEFAULT_LLAMA_CPP_MODEL_VERBOSITY = True class LlamaCPP(CustomLLM): model_url: Optional[str] = Field( description="The URL llama-cpp model to download and use." ) model_path: Optional[str] = Field( description="The path to the llama-cpp model to use." ) temperature: float = Field( default=DEFAULT_TEMPERATURE, description="The temperature to use for sampling.", gte=0.0, lte=1.0, ) max_new_tokens: int = Field( default=DEFAULT_NUM_OUTPUTS, description="The maximum number of tokens to generate.", gt=0, ) context_window: int = Field( default=DEFAULT_CONTEXT_WINDOW, description="The maximum number of context tokens for the model.", gt=0, ) generate_kwargs: Dict[str, Any] = Field( default_factory=dict, description="Kwargs used for generation." ) model_kwargs: Dict[str, Any] = Field( default_factory=dict, description="Kwargs used for model initialization." ) verbose: bool = Field( default=DEFAULT_LLAMA_CPP_MODEL_VERBOSITY, description="Whether to print verbose output.", ) _model: Any = PrivateAttr() def __init__( self, model_url: Optional[str] = None, model_path: Optional[str] = None, temperature: float = DEFAULT_TEMPERATURE, max_new_tokens: int = DEFAULT_NUM_OUTPUTS, context_window: int = DEFAULT_CONTEXT_WINDOW, callback_manager: Optional[CallbackManager] = None, generate_kwargs: Optional[Dict[str, Any]] = None, model_kwargs: Optional[Dict[str, Any]] = None, verbose: bool = DEFAULT_LLAMA_CPP_MODEL_VERBOSITY, system_prompt: Optional[str] = None, messages_to_prompt: Optional[Callable[[Sequence[ChatMessage]], str]] = None, completion_to_prompt: Optional[Callable[[str], str]] = None, pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT, output_parser: Optional[BaseOutputParser] = None, ) -> None: try: from llama_cpp import Llama except ImportError: raise ImportError( "Could not import llama_cpp library." "Please install llama_cpp with `pip install llama-cpp-python`." "See the full installation guide for GPU support at " "`https://github.com/abetlen/llama-cpp-python`" ) model_kwargs = { **{"n_ctx": context_window, "verbose": verbose}, **(model_kwargs or {}), # Override defaults via model_kwargs } # check if model is cached if model_path is not None: if not os.path.exists(model_path): raise ValueError( "Provided model path does not exist. " "Please check the path or provide a model_url to download." ) else: self._model = Llama(model_path=model_path, **model_kwargs) else: cache_dir = get_cache_dir() model_url = model_url or self._get_model_path_for_version() model_name = os.path.basename(model_url) model_path = os.path.join(cache_dir, "models", model_name) if not os.path.exists(model_path): os.makedirs(os.path.dirname(model_path), exist_ok=True) self._download_url(model_url, model_path) assert os.path.exists(model_path) self._model = Llama(model_path=model_path, **model_kwargs) model_path = model_path generate_kwargs = generate_kwargs or {} generate_kwargs.update( {"temperature": temperature, "max_tokens": max_new_tokens} ) super().__init__( model_path=model_path, model_url=model_url, temperature=temperature, context_window=context_window, max_new_tokens=max_new_tokens, callback_manager=callback_manager, generate_kwargs=generate_kwargs, model_kwargs=model_kwargs, verbose=verbose, system_prompt=system_prompt, messages_to_prompt=messages_to_prompt, completion_to_prompt=completion_to_prompt, pydantic_program_mode=pydantic_program_mode, output_parser=output_parser, ) @classmethod def class_name(cls) -> str: return "LlamaCPP_llm" @property def metadata(self) -> LLMMetadata: """LLM metadata.""" return LLMMetadata( context_window=self._model.context_params.n_ctx, num_output=self.max_new_tokens, model_name=self.model_path, ) def _get_model_path_for_version(self) -> str: """Get model path for the current llama-cpp version.""" import pkg_resources version = pkg_resources.get_distribution("llama-cpp-python").version major, minor, patch = version.split(".") # NOTE: llama-cpp-python<=0.1.78 supports GGML, newer support GGUF if int(major) <= 0 and int(minor) <= 1 and int(patch) <= 78: return DEFAULT_LLAMA_CPP_GGML_MODEL else: return DEFAULT_LLAMA_CPP_GGUF_MODEL def _download_url(self, model_url: str, model_path: str) -> None: completed = False try: print("Downloading url", model_url, "to path", model_path) with requests.get(model_url, stream=True) as r: with open(model_path, "wb") as file: total_size = int(r.headers.get("Content-Length") or "0") if total_size < 1000 * 1000: raise ValueError( "Content should be at least 1 MB, but is only", r.headers.get("Content-Length"), "bytes", ) print("total size (MB):", round(total_size / 1000 / 1000, 2)) chunk_size = 1024 * 1024 # 1 MB for chunk in tqdm( r.iter_content(chunk_size=chunk_size), total=int(total_size / chunk_size), ): file.write(chunk) completed = True except Exception as e: print("Error downloading model:", e) finally: if not completed: print("Download incomplete.", "Removing partially downloaded file.") os.remove(model_path) raise ValueError("Download incomplete.") @llm_chat_callback() def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse: prompt = self.messages_to_prompt(messages) completion_response = self.complete(prompt, formatted=True, **kwargs) return completion_response_to_chat_response(completion_response) @llm_chat_callback() def stream_chat( self, messages: Sequence[ChatMessage], **kwargs: Any ) -> ChatResponseGen: prompt = self.messages_to_prompt(messages) completion_response = self.stream_complete(prompt, formatted=True, **kwargs) return stream_completion_response_to_chat_response(completion_response) @llm_completion_callback() def complete( self, prompt: str, formatted: bool = False, **kwargs: Any ) -> CompletionResponse: self.generate_kwargs.update({"stream": False}) if not formatted: prompt = self.completion_to_prompt(prompt) response = self._model(prompt=prompt, **self.generate_kwargs) return CompletionResponse(text=response["choices"][0]["text"], raw=response) @llm_completion_callback() def stream_complete( self, prompt: str, formatted: bool = False, **kwargs: Any ) -> CompletionResponseGen: self.generate_kwargs.update({"stream": True}) if not formatted: prompt = self.completion_to_prompt(prompt) response_iter = self._model(prompt=prompt, **self.generate_kwargs) def gen() -> CompletionResponseGen: text = "" for response in response_iter: delta = response["choices"][0]["text"] text += delta yield CompletionResponse(delta=delta, text=text, raw=response) return gen()