WRAP Models

koalpaca.py

@@ -1,8 +1,9 @@
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, GenerationConfig
 from peft import PeftModel, PeftConfig
+from model import Model
 
-class KoAlpaca:
+class KoAlpaca(Model):
     def __init__(self):
         peft_model_id = "4n3mone/Komuchat-koalpaca-polyglot-12.8B"
         config = PeftConfig.from_pretrained(peft_model_id)

model.py

@@ -0,0 +1,10 @@
+class Model:
+    def __init__(self):
+        self.model = None
+        self.tokenizer = None
+        self.gen_config = None
+        self.INPUT_FORMAT = None
+    
+    def generate(self, inputs:str) -> str:
+        outputs = inputs
+        return outputs

t5.py

@@ -1,6 +1,6 @@
 from transformers import T5TokenizerFast, T5ForConditionalGeneration, GenerationConfig
-
-class T5:
+from model import Model
+class T5(Model):
     def __init__(self,
             model_dir:str='./models/pko_t5_COMU_patience10',
             max_input_length:int=64,