Messing with fastAPI.

app.py

@@ -4,6 +4,7 @@ import gradio
 import json
 import onnxruntime
 import time
+from datetime import datetime
 from transformers import pipeline
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
@@ -32,7 +33,8 @@ print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 model_name = "MoritzLaurer/deberta-v3-base-zeroshot-v2.0"
 tokenizer_name = "MoritzLaurer/deberta-v3-base-zeroshot-v2.0"
 
-classifier = pipeline(task="zero-shot-classification", model=model_name, tokenizer=tokenizer_name, device="cuda:0")
+classifier_cpu = pipeline(task="zero-shot-classification", model=model_name, tokenizer=tokenizer_name)
+classifier_gpu = pipeline(task="zero-shot-classification", model=model_name, tokenizer=tokenizer_name, device="cuda:0")
 # classifier = pipeline(task="zero-shot-classification", model=model_name, tokenizer=tokenizer_name)
 
 def classify(data_string, request: gradio.Request):
@@ -42,19 +44,27 @@ def classify(data_string, request: gradio.Request):
     data = json.loads(data_string)
 
     # Prevent batch suggestion warning in log.
-    classifier.call_count = 0
+    classifier_cpu.call_count = 0
+    classifier_gpu.call_count = 0
 
     # if 'task' in data and data['task'] == 'few_shot_classification':
     #     return few_shot_classification(data)
     # else:
     start_time = time.time()
-    result = zero_shot_classification(data)
-    print(f"classification took {time.time() - start_time}.")
+    result = {}
+    if (data['cpu'])
+        result = zero_shot_classification_cpu(data)
+    else
+        result = zero_shot_classification_gpu(data)
+    print(f"Classification @ [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] took {time.time() - start_time}.")
     return json.dumps(result)
 
+def zero_shot_classification_cpu(data):
+    return classifier_cpu(data['sequence'], candidate_labels=data['candidate_labels'], hypothesis_template=data['hypothesis_template'], multi_label=data['multi_label'])
+
 @spaces.GPU(duration=3)
-def zero_shot_classification(data):
-    return classifier(data['sequence'], candidate_labels=data['candidate_labels'], hypothesis_template=data['hypothesis_template'], multi_label=data['multi_label'])
+def zero_shot_classification_gpu(data):
+    return classifier_gpu(data['sequence'], candidate_labels=data['candidate_labels'], hypothesis_template=data['hypothesis_template'], multi_label=data['multi_label'])
 
 def create_sequences(data):
     return [data['sequence'] + '\n' + data['hypothesis_template'].format(label) for label in data['candidate_labels']]