updates for phi

handler.py

@@ -6,22 +6,31 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 class EndpointHandler:
     def __init__(self, path=""):
         # load the model
-        tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-128k-instruct", trust_remote_code=True)
+        tokenizer = tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-128k-instruct")
         model = AutoModelForCausalLM.from_pretrained(
             "microsoft/Phi-3-mini-128k-instruct",
-            trust_remote_code=True
-        )
+            device_map="cuda",
+            torch_dtype="auto",
+            trust_remote_code=True,
+        ) 
         # create inference pipeline
         self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         inputs = data.pop("inputs", data)
-
-        for key in ['stop_sequences', 'watermark', 'stop']:
-            if key in inputs:
-                del inputs[key]
-
         parameters = data.pop("parameters", None)
+            
+        # Print parameters for debugging
+        print("Parameters before cleaning:", parameters)
+        
+        # Remove unwanted keys from parameters
+        if parameters is not None:
+            for key in ['stop_sequences', 'watermark', 'stop']:
+                if key in parameters:
+                    del parameters[key]
+
+        # Print parameters after cleaning
+        print("Parameters after cleaning:", parameters)
 
         # pass inputs with all kwargs in data
         if parameters is not None:

requirements.txt

@@ -1 +1,4 @@
-flash-attn==latest
+flash_attn==2.5.8
+torch==2.3.1
+accelerate==0.31.0
+transformers==4.41.2