Completed Prototype

Gemma.py

@@ -1,4 +1,5 @@
-from transformers import AutoTokenizer, Gemma3ForCausalLM
+# from transformers import AutoTokenizer, Gemma3ForCausalLM
+from transformers import pipeline
 from huggingface_hub import login
 import spaces
 import torch
@@ -14,30 +15,32 @@ class GemmaLLM:
         
         model_id = "google/gemma-3-1b-it"
         
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
-        self.model = Gemma3ForCausalLM.from_pretrained(
-            model_id,
-            device_map="cuda" if torch.cuda.is_available() else "cpu",
-            torch_dtype=torch.float16,
-        ).eval()
+        # self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        # self.model = Gemma3ForCausalLM.from_pretrained(
+        #     model_id,
+        #     device_map="cuda" if torch.cuda.is_available() else "cpu",
+        #     torch_dtype=torch.float16,
+        # ).eval()
+        
+        self.model = pipeline("text-generation", model=model_id, torch_dtype=torch.bfloat16, device="cuda")
         
-        self.model = self.model.bfloat16()
-
     @spaces.GPU
     def generate(self, message) -> str:
-        inputs = self.tokenizer.apply_chat_template(
-            message,
-            add_generation_prompt=True,
-            tokenize=True,
-            return_dict=True,
-            return_tensors="pt",
-        ).to(self.model.device)
+        # inputs = self.tokenizer.apply_chat_template(
+        #     message,
+        #     add_generation_prompt=True,
+        #     tokenize=True,
+        #     return_dict=True,
+        #     return_tensors="pt",
+        # ).to(self.model.device)
 
-        input_length = inputs["input_ids"].shape[1]
+        # input_length = inputs["input_ids"].shape[1]
 
-        with torch.inference_mode():
-            outputs = self.model.generate(**inputs, max_new_tokens=1024)[0][input_length:]
-            outputs = self.tokenizer.decode(outputs, skip_special_tokens=True)
+        # with torch.inference_mode():
+        #     outputs = self.model.generate(**inputs, max_new_tokens=1024)[0][input_length:]
+        #     outputs = self.tokenizer.decode(outputs, skip_special_tokens=True)
+        
+        outputs = self.model(message, max_new_tokens=1024)[0]["generated_text"]
         
         return outputs