# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['DEFAULT_MODEL_URL', 'DEFAULT_MODEL_NAME', 'DEFAULT_EMBEDDING_MODEL', 'LLM']

# %% ../nbs/00_core.ipynb 3
from langchain.chains import RetrievalQA
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.vectorstores import Chroma
from langchain.llms import GPT4All, LlamaCpp
import chromadb
import os
import argparse
import time
import warnings
import sys
from typing import Any, Dict, Generator, List, Optional, Tuple, Union


# %% ../nbs/00_core.ipynb 4
from . import utils as U
DEFAULT_MODEL_URL = 'https://huggingface.co/TheBloke/Wizard-Vicuna-7B-Uncensored-GGML/resolve/main/Wizard-Vicuna-7B-Uncensored.ggmlv3.q4_0.bin'
DEFAULT_MODEL_NAME = os.path.basename(DEFAULT_MODEL_URL)
DEFAULT_EMBEDDING_MODEL = 'sentence-transformers/all-MiniLM-L6-v2'

class LLM:
    def __init__(self, 
                 model_name=DEFAULT_MODEL_NAME,
                 n_gpu_layers:Optional[int]=None, 
                 max_tokens:int=512, 
                 n_ctx:int=2048, 
                 n_batch:int=1024,
                 mute_stream=False,
                verbose=False):
        """
        LLM Constructor
        
        **Args:**
        
        - *n_gpu_layers*: Number of layers to be loaded into gpu memory. Default is `None`.
        - *max_tokens*: The maximum number of tokens to generate.
        - *n_ctx*: Token context window.
        - *n_batch*: Number of tokens to process in parallel.
        - *mute_stream*: Mute ChatGPT-like token stream output during generation
        - *verbose*: Verbosity
        """
        self.model_name = model_name
        if not os.path.isfile(os.path.join(U.get_datadir(), model_name)):
            warnings.warn('The model {model_name} does not exist in {U.get_datadir()}. '+\
                          'Please execute LLM.download() to download it.')
        self.llm = None
        self.n_gpu_layers = n_gpu_layers
        self.max_tokens = max_tokens
        self.n_ctx = n_ctx
        self.n_batch = n_batch
        self.callbacks = [] if mute_stream else [StreamingStdOutCallbackHandler()]
        self.verbose = verbose
 
    @classmethod
    def download_model(cls, model_url=DEFAULT_MODEL_NAME, confirm=True):
        """
        Download an LLM in GGML format supported by [lLama.cpp](https://github.com/ggerganov/llama.cpp).
        """
        datadir = U.get_datadir()
        model_name = os.path.basename(model_url)
        filename = os.path.join(datadir, model_name)
        confirm_msg = f"You are about to download the LLM {model_name} to the {datadir} folder. Are you sure?"
        if os.path.isfile(filename):
            confirm_msg = f'There is already a file {model_name} in {datadir}. Do you want to still download it?'
            
        shall = True
        if confirm:
            shall = input("%s (Y/n) " % confirm_msg) == "Y"
        if shall:
            U.download(model_url, filename)
        else:
            warnings.warn(f'{model_name} was not downloaded because "Y" was not selected.')
        return
    
    def ingest(self, 
               source_directory:str,
               embedding_model_name:str ='sentence-transformers/all-MiniLM-L6-v2',
               embedding_model_kwargs:dict ={'device': 'cpu'}
              ):
        """
        Ingests all documents in `source_folder` into vector database.
        Previously-ingested documents are ignored.

        **Args:**
        
        - *source_directory*: path to folder containing document store
        - *embedding_model*: name of sentence-transformers model
        - *embedding_model_kwargs*: arguments to embedding model (e.g., `{device':'cpu'}`)
        
        **Returns:** `None`
        """
        import onprem.ingest as ingest_fn
        ingest_fn(source_directory, 
                  embedding_model_name=embedding_model_name,
                  embedding_model_kwargs=embedding_model_kwargs)
 
        
    def check_model(self):
        datadir = U.get_datadir()
        model_path = os.path.join(datadir, self.model_name)
        if not os.path.isfile(model_path):
            raise ValueError(f'The LLM model {self.model_name} does not appear to have been downloaded. '+\
                             f'Execute the download_model() method to download it.')
        return model_path
        
        
    def prompt(self, prompt):
        """
        Send prompt to LLM to generate a response
        """
        model_path = self.check_model()
        
        if not self.llm:
            self.llm =  llm = LlamaCpp(model_path=model_path, 
                                       max_tokens=self.max_tokens, 
                                       n_batch=self.n_batch, 
                                       callbacks=self.callbacks, 
                                       verbose=self.verbose, 
                                       n_gpu_layers=self.n_gpu_layers, 
                                       n_ctx=self.n_ctx)
        return self.llm(prompt)    
