Pavankalyan commited on
Commit
8f4c005
1 Parent(s): df64ef2

Create Dockerfile

Browse files
Files changed (1) hide show
  1. Dockerfile +34 -0
Dockerfile ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM python:3.9-slim-bullseye
2
+
3
+ RUN apt-get -y update && \
4
+ apt-get install -y --no-install-recommends build-essential \
5
+ curl wget nginx ca-certificates npm \
6
+ && npm install pm2 -g \
7
+ && pip install --upgrade pip setuptools \
8
+ && rm -rf /var/lib/apt/lists/*
9
+
10
+ COPY requirements.txt .
11
+ RUN pip install -r requirements.txt
12
+
13
+ class ZeroShotTextClassifier:
14
+ # Class variable for the model
15
+ classifier = None
16
+ @classmethod
17
+ def load(cls):
18
+ if cls.classifier is None:
19
+ # Load the model only once
20
+ cls.classifier = pipeline("zero-shot-classification",
21
+ model="facebook/bart-large-mnli")
22
+ @classmethod
23
+ def predict(cls, text, candidate_labels):
24
+ # Ensure the model is loaded
25
+ cls.load()
26
+ # Predict
27
+ huggingface_predictions = cls.classifier(text, candidate_labels)
28
+ # Create our own prediction object with the best label
29
+ max_index = np.argmax(huggingface_predictions["scores"])
30
+ label = huggingface_predictions["labels"][max_index]
31
+ score = huggingface_predictions["scores"][max_index]
32
+ return {"label": label, "score": score}
33
+
34
+ RUN python -c "from transformers import pipeline; classifier = pipeline('zero-shot-classification', model='facebook/bart-large-mnli')"