Update handler.py

handler.py

@@ -17,9 +17,10 @@ class EndpointHandler:
         print(f"Loading model from {path} on device: {self.device}...")
         self.model = AutoModelForCausalLM.from_pretrained(
             path,
-            torch_dtype=torch.float16,
+            torch_dtype=torch.bfloat16,     # ✅ Changed to bfloat16 
             trust_remote_code=True,
             device_map="auto",
+            attn_implementation="flash_attention_2"  # ✅ Faster + stable
         )
         self.model.eval()
         print("✅ Model loaded successfully!")
@@ -31,6 +32,9 @@ class EndpointHandler:
         if not isinstance(inputs, str) or not inputs.strip():
             return {"generated_text": ""}
 
+         # ✅ StarCoder2: Add code context prefix
+        prompt = f"<fim_prefix>{inputs}<fim_suffix><fim_middle>"
+
         gen_kwargs = {
             "max_new_tokens": min(parameters.get("max_new_tokens", 256), 512),  # Cap for stability
             "temperature": parameters.get("temperature", 0.2),
@@ -42,11 +46,13 @@ class EndpointHandler:
             "pad_token_id": self.tokenizer.pad_token_id,
         }
         
-        print(f"Generating with parameters: {gen_kwargs}")
+        # print(f"Generating with parameters: {gen_kwargs}")
+         print(f"Prompt length: {len(prompt)} | Gen params: {gen_kwargs}")
         
         # StarCoder2 tokenization
         inputs = inputs.strip()
         tokenized = self.tokenizer(
+            prompt,
             inputs,
             return_tensors="pt",
             truncation=True,
@@ -65,10 +71,11 @@ class EndpointHandler:
             
             # Extract ONLY newly generated tokens
             new_tokens = outputs[0][len(tokenized.input_ids[0]):]
-            generated_text = self.tokenizer.decode(
+            generated = self.tokenizer.decode(
                 new_tokens, 
                 skip_special_tokens=True,
                 clean_up_tokenization_spaces=True
             )
         
+            generated = generated.replace("<fim_middle>", "").replace("<fim_suffix>", "").strip()
         return {"generated_text": generated_text.strip()}