Zwea Htet commited on
Commit
5ef1e14
·
1 Parent(s): 9281dcc

update llamacustom

Browse files
Files changed (1) hide show
  1. models/llamaCustom.py +5 -22
models/llamaCustom.py CHANGED
@@ -66,25 +66,10 @@ def load_model(model_name: str):
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:
@@ -138,9 +123,7 @@ class LlamaCustom:
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(
 
66
 
67
 
68
  class OurLLM(CustomLLM):
69
+ def __init__(self, model_name: str, pipeline):
70
+ super().__init__() # Call the __init__ method of CustomLLM
71
+ self.model_name = model_name
72
+ self.pipeline = pipeline
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
73
 
74
  @property
75
  def metadata(self) -> LLMMetadata:
 
123
 
124
  # define llm
125
  pipe = load_model(model_name=model_name)
126
+ llm = OurLLM(model_name=model_name, pipeline=pipe)
 
 
127
 
128
  llm_predictor = LLMPredictor(llm=llm)
129
  service_context = ServiceContext.from_defaults(