changed class interface with iterator

interface.py

@@ -9,15 +9,15 @@ from pydantic import Field, field_validator
 
 # for transformers 2
 class GemmaLLMInterface(CustomLLM):
-    def __init__(self, model_name: str = "google/gemma-2b-it", **kwargs):
+    def __init__(self, model_id: str = "google/gemma-2-2b-it", **kwargs):
         super().__init__(**kwargs)
-        self.model_name = model_name
+        self.model_id = model_id
         self.model = AutoModelForCausalLM.from_pretrained(
-            self.model_name,
+            self.model_id,
             device_map="auto",
             torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
         )
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
         self.context_window = 8192
         self.num_output = 2048
 
@@ -32,7 +32,7 @@ class GemmaLLMInterface(CustomLLM):
         return LLMMetadata(
             context_window=self.context_window,
             num_output=self.num_output,
-            model_name=self.model_name,
+            model_name=self.model_id,
         )
 
     @llm_completion_callback()