added new class

backend.py

@@ -20,23 +20,23 @@ login(huggingface_token)
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-"""model_id = "google/gemma-2-2b-it"
+model_id = "google/gemma-2-2b-it"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     device_map="auto", ## change this back to auto!!!
     torch_dtype= torch.bfloat16 if torch.cuda.is_available() else torch.float32,
     token=True)
-model.eval()"""
+model.eval()
 
 #from accelerate import disk_offload
 #disk_offload(model=model, offload_dir="offload")
 
 # what models will be used by LlamaIndex:
 """Settings.embed_model = InstructorEmbedding(model_name="hkunlp/instructor-base")
-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")
+Settings.llm  = GemmaLLMInterface(model=model)"""
+gemma_llm = GemmaLLMInterface(model_name=model_id)
+
 
 ############################---------------------------------
 

interface.py

@@ -9,66 +9,48 @@ from pydantic import Field, field_validator
 
 # for transformers 2
 class GemmaLLMInterface(CustomLLM):
-    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)
+    model: AutoModelForCausalLM = None
+    tokenizer: AutoTokenizer = None
+    context_window: int = 8192
+    num_output: int = 2048
+    model_name: str = "gemma_2"
+
+    def __init__(self, model_name: str):
         self.model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            device_map="auto",
-            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+            model_name,
+            device_map="auto",  # Set device mapping automatically
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
         )
-        self.model.eval()
-    
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+
     def _format_prompt(self, message: str) -> str:
-        return f"<start_of_turn>user\n{message}<end_of_turn>\n<start_of_turn>model\n"
+        return (
+            f"<start_of_turn>user\n{message}<end_of_turn>\n" f"<start_of_turn>model\n"
+        )
 
     @property
     def metadata(self) -> LLMMetadata:
+        """Get LLM metadata."""
         return LLMMetadata(
             context_window=self.context_window,
             num_output=self.num_output,
-            model_name=self.model_id,
+            model_name=self.model_name,
         )
-    
-    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:
-        inputs = self._prepare_inputs(prompt)
-        response = "".join(self._generate(inputs))
+        prompt = self._format_prompt(prompt)
+        inputs = self.tokenizer(prompt, return_tensors="pt")  # Tokenize the prompt
+        raw_response = self.model.generate(**inputs, max_length=self.num_output)
+        response = self.tokenizer.decode(raw_response[0], skip_special_tokens=True)
         return CompletionResponse(text=response)
 
     @llm_completion_callback()
     def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
-        inputs = self._prepare_inputs(prompt)
-        response = ""
-        for new_token in self._generate(inputs):
-            response += new_token
-            yield CompletionResponse(text=response, delta=new_token)
+        response = self.complete(prompt).text
+        for token in response:
+            response += token
+            yield CompletionResponse(text=response, delta=token)
 
 
 # for transformers 1