zhtet commited on
Commit
ad2d0d7
·
1 Parent(s): 21fbffd

Update models/llamaCustom.py

Browse files
Files changed (1) hide show
  1. models/llamaCustom.py +12 -11
models/llamaCustom.py CHANGED
@@ -73,16 +73,6 @@ class CustomLLM(LLM):
73
 
74
  @st.cache_resource
75
  class LlamaCustom:
76
- # define llm
77
- prompt_helper = PromptHelper(
78
- context_window=CONTEXT_WINDOW,
79
- num_output=NUM_OUTPUT,
80
- chunk_overlap_ratio=CHUNK_OVERLAP_RATION,
81
- )
82
- llm_predictor = LLMPredictor(llm=CustomLLM())
83
- service_context = ServiceContext.from_defaults(
84
- llm_predictor=llm_predictor, prompt_helper=prompt_helper
85
- )
86
 
87
  def __init__(self, model_name: str) -> None:
88
  self.vector_index = self.initialize_index(model_name=model_name)
@@ -103,11 +93,22 @@ class LlamaCustom:
103
  # index = pickle.loads(file.readlines())
104
  return index
105
  else:
 
 
 
 
 
 
 
 
 
 
 
106
  # documents = prepare_data(r"./assets/regItems.json")
107
  documents = SimpleDirectoryReader(input_dir="./assets/pdf").load_data()
108
 
109
  index = GPTVectorStoreIndex.from_documents(
110
- documents, service_context=self.service_context
111
  )
112
 
113
  # local write access
 
73
 
74
  @st.cache_resource
75
  class LlamaCustom:
 
 
 
 
 
 
 
 
 
 
76
 
77
  def __init__(self, model_name: str) -> None:
78
  self.vector_index = self.initialize_index(model_name=model_name)
 
93
  # index = pickle.loads(file.readlines())
94
  return index
95
  else:
96
+ # define llm
97
+ prompt_helper = PromptHelper(
98
+ context_window=CONTEXT_WINDOW,
99
+ num_output=NUM_OUTPUT,
100
+ chunk_overlap_ratio=CHUNK_OVERLAP_RATION,
101
+ )
102
+ llm_predictor = LLMPredictor(llm=CustomLLM())
103
+ service_context = ServiceContext.from_defaults(
104
+ llm_predictor=llm_predictor, prompt_helper=prompt_helper
105
+ )
106
+
107
  # documents = prepare_data(r"./assets/regItems.json")
108
  documents = SimpleDirectoryReader(input_dir="./assets/pdf").load_data()
109
 
110
  index = GPTVectorStoreIndex.from_documents(
111
+ documents, service_context=service_context
112
  )
113
 
114
  # local write access