File size: 763 Bytes
90653e1
 
 
 
 
 
9f10a81
90653e1
d15a232
 
93b834c
90653e1
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import os

from dotenv import load_dotenv
from llama_index import (Document, GPTSimpleVectorIndex, LLMPredictor,
                         ServiceContext)

from data.prepare import data

from .customLLM import CustomLLM, prompt_helper

load_dotenv()

#define our llm
llm_predictor = LLMPredictor(llm=CustomLLM())
service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)

def initialize_index(index_name):
    if os.path.exists(index_name):
        return GPTSimpleVectorIndex.load_from_disk(index_name)
    else:
        documents = [Document(d) for d in data]
        index = GPTSimpleVectorIndex.from_documents(documents, service_context=service_context)
        index.save_to_disk(index_name)
        return index