Spaces:
Sleeping
Sleeping
Zwea Htet
commited on
Commit
·
71b7421
1
Parent(s):
8b30acb
updated custom class
Browse files- models/llamaCustom.py +5 -0
models/llamaCustom.py
CHANGED
|
@@ -59,7 +59,12 @@ def load_model(mode_name: str):
|
|
| 59 |
|
| 60 |
|
| 61 |
class CustomLLM(LLM):
|
|
|
|
|
|
|
|
|
|
| 62 |
def __init__(self, model_name: str):
|
|
|
|
|
|
|
| 63 |
self.llm_model_name = model_name
|
| 64 |
self.pipeline = load_model(mode_name=model_name)
|
| 65 |
|
|
|
|
| 59 |
|
| 60 |
|
| 61 |
class CustomLLM(LLM):
|
| 62 |
+
llm_model_name: str
|
| 63 |
+
pipeline: Any
|
| 64 |
+
|
| 65 |
def __init__(self, model_name: str):
|
| 66 |
+
super().__init__()
|
| 67 |
+
|
| 68 |
self.llm_model_name = model_name
|
| 69 |
self.pipeline = load_model(mode_name=model_name)
|
| 70 |
|