import torch
from langchain.llms.base import LLM
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from llama_index import LangchainEmbedding
from llama_index import SimpleDirectoryReader, LLMPredictor, PromptHelper, GPTSimpleVectorIndex
from peft import PeftModel
from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig


hf_model_path = "models/llama-7b"
alpaca_model_path = "models/lora-alpaca"

tokenizer = LlamaTokenizer.from_pretrained(hf_model_path)

model = LlamaForCausalLM.from_pretrained(
    hf_model_path,
    load_in_8bit=True, #Dissabling could solve some errors
    device_map="auto",
)
model = PeftModel.from_pretrained(model, alpaca_model_path)


device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
max_length = 1500 #2048
max_new_tokens = 48


class LLaMALLM(LLM):
    def _call(self, prompt, stop=None):
        prompt += "### Response:"

        inputs = tokenizer(prompt, return_tensors="pt")
        input_ids = inputs["input_ids"].cuda()
        
        generation_config = GenerationConfig(
            temperature=0.6,
            top_p=0.95,
            repetition_penalty=1.15,
        )
        with torch.no_grad():
            generation_output = model.generate(
                input_ids=input_ids,
                generation_config=generation_config,
                return_dict_in_generate=True,
                output_scores=True,
                max_new_tokens=128,
            )
        response = ""
        for s in generation_output.sequences:
            response += tokenizer.decode(s)
            
        response = response[len(prompt):]
        print("Model Response:", response)
        return response

    def _identifying_params(self):
        return {"name_of_model": "alpaca"}

    def _llm_type(self):
        return "custom"

max_input_size = max_length
num_output = max_new_tokens
max_chunk_overlap = 20

prompt_helper = PromptHelper(max_input_size, num_output, max_chunk_overlap)
embed_model = LangchainEmbedding(HuggingFaceEmbeddings())
documents = SimpleDirectoryReader('data').load_data()
llm_predictor = LLMPredictor(llm=LLaMALLM())
index = GPTSimpleVectorIndex(documents, llm_predictor=llm_predictor, embed_model=embed_model, prompt_helper=prompt_helper)

index.save_to_disk('index.json')
new_index = GPTSimpleVectorIndex.load_from_disk('index.json', embed_model=embed_model, llm_predictor=llm_predictor, prompt_helper=prompt_helper)

response = new_index.query("What did Gatsby do before he met Daisy?")
print(response.response)

response = new_index.query("What did the narrator do after getting back to Chicago?")
print(response.response)

langchain.embeddings.huggingface.HuggingFaceEmbeddings
[((460, 505), 'transformers.LlamaTokenizer.from_pretrained', 'LlamaTokenizer.from_pretrained', (['hf_model_path'], {}), '(hf_model_path)\n', (490, 505), False, 'from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig\n'), ((517, 606), 'transformers.LlamaForCausalLM.from_pretrained', 'LlamaForCausalLM.from_pretrained', (['hf_model_path'], {'load_in_8bit': '(True)', 'device_map': '"""auto"""'}), "(hf_model_path, load_in_8bit=True,\n    device_map='auto')\n", (549, 606), False, 'from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig\n'), ((667, 718), 'peft.PeftModel.from_pretrained', 'PeftModel.from_pretrained', (['model', 'alpaca_model_path'], {}), '(model, alpaca_model_path)\n', (692, 718), False, 'from peft import PeftModel\n'), ((1997, 2056), 'llama_index.PromptHelper', 'PromptHelper', (['max_input_size', 'num_output', 'max_chunk_overlap'], {}), '(max_input_size, num_output, max_chunk_overlap)\n', (2009, 2056), False, 'from llama_index import SimpleDirectoryReader, LLMPredictor, PromptHelper, GPTSimpleVectorIndex\n'), ((2226, 2345), 'llama_index.GPTSimpleVectorIndex', 'GPTSimpleVectorIndex', (['documents'], {'llm_predictor': 'llm_predictor', 'embed_model': 'embed_model', 'prompt_helper': 'prompt_helper'}), '(documents, llm_predictor=llm_predictor, embed_model=\n    embed_model, prompt_helper=prompt_helper)\n', (2246, 2345), False, 'from llama_index import SimpleDirectoryReader, LLMPredictor, PromptHelper, GPTSimpleVectorIndex\n'), ((2390, 2526), 'llama_index.GPTSimpleVectorIndex.load_from_disk', 'GPTSimpleVectorIndex.load_from_disk', (['"""index.json"""'], {'embed_model': 'embed_model', 'llm_predictor': 'llm_predictor', 'prompt_helper': 'prompt_helper'}), "('index.json', embed_model=embed_model,\n    llm_predictor=llm_predictor, prompt_helper=prompt_helper)\n", (2425, 2526), False, 'from llama_index import SimpleDirectoryReader, LLMPredictor, PromptHelper, GPTSimpleVectorIndex\n'), ((757, 782), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (780, 782), False, 'import torch\n'), ((733, 753), 'torch.device', 'torch.device', (['"""cuda"""'], {}), "('cuda')\n", (745, 753), False, 'import torch\n'), ((2091, 2114), 'langchain.embeddings.huggingface.HuggingFaceEmbeddings', 'HuggingFaceEmbeddings', ([], {}), '()\n', (2112, 2114), False, 'from langchain.embeddings.huggingface import HuggingFaceEmbeddings\n'), ((1088, 1158), 'transformers.GenerationConfig', 'GenerationConfig', ([], {'temperature': '(0.6)', 'top_p': '(0.95)', 'repetition_penalty': '(1.15)'}), '(temperature=0.6, top_p=0.95, repetition_penalty=1.15)\n', (1104, 1158), False, 'from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig\n'), ((2129, 2158), 'llama_index.SimpleDirectoryReader', 'SimpleDirectoryReader', (['"""data"""'], {}), "('data')\n", (2150, 2158), False, 'from llama_index import SimpleDirectoryReader, LLMPredictor, PromptHelper, GPTSimpleVectorIndex\n'), ((1224, 1239), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (1237, 1239), False, 'import torch\n')]