初始化项目,由ModelHub XC社区提供模型
Model: hyperspaceai/hyperEngine_phi3_128k Source: Original Platform
This commit is contained in:
25
handler.py
Normal file
25
handler.py
Normal file
@@ -0,0 +1,25 @@
|
||||
import torch
|
||||
from typing import Dict, List, Any
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
|
||||
|
||||
|
||||
class EndpointHandler():
|
||||
def __init__(self, path=""):
|
||||
model = AutoModelForCausalLM.from_pretrained("hyperspaceai/hyperEngine_phi3_128k", device_map="auto", torch_dtype="auto", trust_remote_code=True)
|
||||
tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-128k-instruct")
|
||||
self.pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
|
||||
|
||||
def __call__(self, data:Dict[str, Any]) :
|
||||
messages = data.pop("messages", None)
|
||||
generation_args = data.pop("generation_args", None)
|
||||
|
||||
if generation_args==None :
|
||||
generation_args = {
|
||||
"max_new_tokens": 500,
|
||||
"return_full_text": False,
|
||||
"temperature": 0.0,
|
||||
"do_sample": False,
|
||||
}
|
||||
|
||||
output = self.pipe(messages, **generation_args)
|
||||
return output[0]['generated_text']
|
||||
Reference in New Issue
Block a user