Zwea Htet commited on
Commit
7704dca
·
1 Parent(s): 5d94ab8

updated code

Browse files
Files changed (1) hide show
  1. models/llamaCustom.py +23 -5
models/llamaCustom.py CHANGED
@@ -66,9 +66,25 @@ def load_model(model_name: str):
66
 
67
 
68
  class OurLLM(CustomLLM):
69
- def __init__(self, model_name: str, model_pipeline):
70
- self.model_name = model_name
71
- self.pipeline = model_pipeline
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
72
 
73
  @property
74
  def metadata(self) -> LLMMetadata:
@@ -76,7 +92,7 @@ class OurLLM(CustomLLM):
76
  return LLMMetadata(
77
  context_window=CONTEXT_WINDOW,
78
  num_output=NUM_OUTPUT,
79
- model_name=self.model_name,
80
  )
81
 
82
  @llm_completion_callback()
@@ -122,7 +138,9 @@ class LlamaCustom:
122
 
123
  # define llm
124
  pipe = load_model(model_name=model_name)
125
- llm = OurLLM(model_name=model_name, model_pipeline=pipe)
 
 
126
 
127
  llm_predictor = LLMPredictor(llm=llm)
128
  service_context = ServiceContext.from_defaults(
 
66
 
67
 
68
  class OurLLM(CustomLLM):
69
+ def __init__(self):
70
+ self.model_name = ""
71
+ self.pipeline = None
72
+
73
+ @property
74
+ def modelName(self):
75
+ return self.model_name
76
+
77
+ @property
78
+ def pipeline(self):
79
+ return self.pipeline
80
+
81
+ @modelName.setter
82
+ def modelName(self, name: str):
83
+ self.model_name = name
84
+
85
+ @pipeline.setter
86
+ def pipeline(self, value):
87
+ self.pipeline = value
88
 
89
  @property
90
  def metadata(self) -> LLMMetadata:
 
92
  return LLMMetadata(
93
  context_window=CONTEXT_WINDOW,
94
  num_output=NUM_OUTPUT,
95
+ # model_name=self.model_name,
96
  )
97
 
98
  @llm_completion_callback()
 
138
 
139
  # define llm
140
  pipe = load_model(model_name=model_name)
141
+ llm = OurLLM()
142
+ llm.model_name = model_name
143
+ llm.pipeline = pipe
144
 
145
  llm_predictor = LLMPredictor(llm=llm)
146
  service_context = ServiceContext.from_defaults(