added new class

backend.py

@@ -38,7 +38,6 @@ Settings.llm  = GemmaLLMInterface(model=model, tokenizer=tokenizer)"""
 Settings.embed_model = InstructorEmbedding(model_name="hkunlp/instructor-base")
 Settings.llm = GemmaLLMInterface(model_id="google/gemma-2-2b-it")
 
-
 ############################---------------------------------
 
 # Get the parser

interface.py

@@ -8,14 +8,17 @@ from threading import Thread
 
 # for transformers 2
 class GemmaLLMInterface(CustomLLM):
-    def __init__(self, model_id: str = "google/gemma-2-2b-it", context_window: int = 8192, num_output: int = 2048):
-        self.model_id = model_id
-        self.context_window = context_window
-        self.num_output = num_output
-        
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model_id: str = Field(default="google/gemma-2-2b-it")
+    context_window: int = Field(default=8192)
+    num_output: int = Field(default=2048)
+    tokenizer: Any = Field(default=None)
+    model: Any = Field(default=None)
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
         self.model = AutoModelForCausalLM.from_pretrained(
-            model_id,
+            self.model_id,
             device_map="auto",
             torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
         )