syberWolf commited on
Commit
001e799
·
1 Parent(s): 96a0103

change handler

Browse files
Files changed (1) hide show
  1. handler.py +2 -2
handler.py CHANGED
@@ -2,8 +2,8 @@ from typing import Dict, List, Any
2
  from transformers import pipeline
3
 
4
  class EndpointHandler:
5
- def __init__(self, model_name="Qwen/Qwen2-1.5B-Instruct"):
6
- self.pipeline = pipeline("text-generation", model=model_name) # Note: Model name provided as argument for flexibility
7
 
8
  def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
9
  inputs = data.pop("inputs", data)
 
2
  from transformers import pipeline
3
 
4
  class EndpointHandler:
5
+ def __init__(self, path=""):
6
+ self.pipeline = pipeline("text-generation", model="Qwen/Qwen2-1.5B-Instruct")
7
 
8
  def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
9
  inputs = data.pop("inputs", data)