Messing with configuration.

app.py

@@ -1,4 +1,5 @@
 import spaces
+import torch
 import gradio
 import json
 import onnxruntime
@@ -19,8 +20,8 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# print(f"Is CUDA available: {torch.cuda.is_available()}")
-# print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+print(f"Is CUDA available: {torch.cuda.is_available()}")
+print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 
 # "xenova/mobilebert-uncased-mnli" "typeform/mobilebert-uncased-mnli" Fast but small--same as bundled in Statosphere
 # "xenova/deberta-v3-base-tasksource-nli" Not impressed
@@ -36,8 +37,8 @@ tokenizer_name = "MoritzLaurer/deberta-v3-base-zeroshot-v2.0"
 # model = ORTModelForSequenceClassification.from_pretrained(model_name, export=True, provider="CUDAExecutionProvider")
 # tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, model_max_length=512)
 
-# classifier = 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)
+classifier = 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):
     if request:
@@ -49,7 +50,7 @@ def classify(data_string, request: gradio.Request):
     # else:
     return zero_shot_classification(data)
 
-@spaces.GPU
+@spaces.GPU(duration=10)
 def zero_shot_classification(data):
     results = classifier(data['sequence'], candidate_labels=data['candidate_labels'], hypothesis_template=data['hypothesis_template'], multi_label=data['multi_label'])
     response_string = json.dumps(results)

requirements.txt

@@ -3,5 +3,5 @@ fastapi==0.88.0
 huggingface_hub==0.26.0
 json5==0.9.25
 numpy
-optimum[onnxruntime]==1.24.0
+optimum[onnxruntime-gpu]==1.24.0
 transformers==4.36