test qwen

handler.py

@@ -1,4 +1,3 @@
-import torch
 from typing import Dict, List, Any
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
@@ -6,13 +5,12 @@ 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")
+        tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B-Instruct")
         model = AutoModelForCausalLM.from_pretrained(
-            "microsoft/Phi-3-mini-128k-instruct",
-            device_map="cuda",
+            "Qwen/Qwen2-1.5B-Instruct",
             torch_dtype="auto",
-            trust_remote_code=True,
-        ) 
+            device_map="auto"
+        )
         # create inference pipeline
         self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
@@ -26,4 +24,4 @@ class EndpointHandler:
         else:
             prediction = self.pipeline(inputs)
         # postprocess the prediction
-        return prediction
+        return prediction

requirements.txt

@@ -1,4 +0,0 @@
-flash_attn==2.5.8
-torch==2.3.1
-accelerate==0.31.0
-transformers==4.41.2