Wrap model to HPU + update Dockerfile packages

This commit is contained in:
PiotrBLL 2024-11-11 16:52:01 +01:00
parent 50adb7654b
commit f6ccd01567
2 changed files with 11 additions and 4 deletions

View File

@ -42,5 +42,8 @@ COPY requirements.txt /root/whisper/requirements.txt
# Install Python packages from requirements.txt
RUN pip install --upgrade pip \
&& pip install optimum[habana]==1.13.1 \
&& pip install -r requirements.txt
&& pip install optimum-habana==1.14.1 \
&& pip install transformers==4.33.0 \
&& pip install huggingface-hub==0.26.2 \
&& pip install tiktoken==0.8.0 \
&& pip install numba==0.60.0

View File

@ -165,10 +165,14 @@ def load_model(
if device == "hpu":
from habana_frameworks.torch.utils.library_loader import load_habana_module
from habana_frameworks.torch.hpu import wrap_in_hpu_graph
load_habana_module()
if torch.hpu.is_available():
return wrap_in_hpu_graph(model)
from habana_frameworks.torch.hpu import wrap_in_hpu_graph
model = model.eval().to(device)
model = wrap_in_hpu_graph(model)
return model
return model.to(device)