added new class

interface.py

@@ -9,27 +9,18 @@ from pydantic import Field, field_validator
 
 # for transformers 2
 class GemmaLLMInterface(CustomLLM):
-    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)
-
-    # Validators are restructured to avoid deepcopy issues
-    @field_validator('context_window', 'num_output', pre=True)
-    def ensure_integer(cls, v):
-        return int(v) if isinstance(v, str) else v
-
-    def __init__(self, **data):
-        super().__init__(**data)
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+    def __init__(self, model_id: str = "google/gemma-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)
         self.model = AutoModelForCausalLM.from_pretrained(
-            self.model_id,
+            model_id,
             device_map="auto",
             torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
         )
         self.model.eval()
-
+    
     def _format_prompt(self, message: str) -> str:
         return f"<start_of_turn>user\n{message}<end_of_turn>\n<start_of_turn>model\n"
 
@@ -40,42 +31,44 @@ class GemmaLLMInterface(CustomLLM):
             num_output=self.num_output,
             model_name=self.model_id,
         )
-
+    
+    def _prepare_inputs(self, prompt: str) -> dict:
+        formatted_prompt = self._format_prompt(prompt)
+        inputs = self.tokenizer(formatted_prompt, return_tensors="pt", add_special_tokens=True).to(self.model.device)
+        if inputs["input_ids"].shape[1] > self.context_window:
+            inputs["input_ids"] = inputs["input_ids"][:, -self.context_window:]
+        return inputs
+    
+    def _generate(self, inputs: dict) -> Iterator[str]:
+        for output in self.model.generate(
+            **inputs,
+            max_new_tokens=self.num_output,
+            do_sample=True,
+            top_p=0.9,
+            top_k=50,
+            temperature=0.7,
+            num_beams=1,
+            repetition_penalty=1.1,
+            streamer=None,
+            return_dict_in_generate=True,
+            output_scores=False,
+        ):
+            new_tokens = output.sequences[:, inputs["input_ids"].shape[-1]:]
+            yield self.tokenizer.decode(new_tokens[0], skip_special_tokens=True)
+           
     @llm_completion_callback()
     def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
-        formatted_prompt = self._format_prompt(prompt)
-        inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.model.device)
-
-        with torch.no_grad():
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=self.num_output,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.95,
-            )
-
-        response = self.tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+        inputs = self._prepare_inputs(prompt)
+        response = "".join(self._generate(inputs))
         return CompletionResponse(text=response)
 
     @llm_completion_callback()
     def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
-        formatted_prompt = self._format_prompt(prompt)
-        inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.model.device)
-
+        inputs = self._prepare_inputs(prompt)
         response = ""
-        with torch.no_grad():
-            for output in self.model.generate(
-                **inputs,
-                max_new_tokens=self.num_output,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.95,
-                streamer=True,
-            ):
-                token = self.tokenizer.decode(output, skip_special_tokens=True)
-                response += token
-                yield CompletionResponse(text=response, delta=token)
+        for new_token in self._generate(inputs):
+            response += new_token
+            yield CompletionResponse(text=response, delta=new_token)
 
 
 # for transformers 1