Messing with configuration.

app.py

@@ -5,7 +5,6 @@ import torch
 from transformers import AutoTokenizer
 from transformers import pipeline
 from optimum.onnxruntime import ORTModelForSequenceClassification
-from optimum.onnxruntime import ORTModelForFeatureExtraction
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from setfit import SetFitModel, SetFitTrainer, Trainer, TrainingArguments
@@ -27,37 +26,6 @@ app.add_middleware(
 print(f"Is CUDA available: {torch.cuda.is_available()}")
 print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 
-# class OnnxSetFitModel:
-#     def __init__(self, ort_model, tokenizer, model_head):
-#         self.ort_model = ort_model
-#         self.tokenizer = tokenizer
-#         self.model_head = model_head
-
-#     def predict(self, inputs):
-#         encoded_inputs = self.tokenizer(
-#             inputs, padding=True, truncation=True, return_tensors="pt"
-#         ).to(self.ort_model.device)
-
-#         outputs = self.ort_model(**encoded_inputs)
-#         embeddings = mean_pooling(
-#             outputs["last_hidden_state"], encoded_inputs["attention_mask"]
-#         )
-#         return self.model_head.predict(embeddings.cpu())
-
-#     def predict_proba(self, inputs):
-#         encoded_inputs = self.tokenizer(
-#             inputs, padding=True, truncation=True, return_tensors="pt"
-#         ).to(self.ort_model.device)
-
-#         outputs = self.ort_model(**encoded_inputs)
-#         embeddings = mean_pooling(
-#             outputs["last_hidden_state"], encoded_inputs["attention_mask"]
-#         )
-#         return self.model_head.predict_proba(embeddings.cpu())
-
-#     def __call__(self, inputs):
-#         return self.predict(inputs)
-
 # "xenova/mobilebert-uncased-mnli" "typeform/mobilebert-uncased-mnli" Fast but small--same as bundled in Statosphere
 # "xenova/deberta-v3-base-tasksource-nli" Not impressed
 # "Xenova/bart-large-mnli" A bit slow
@@ -67,34 +35,9 @@ print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 model_name = "MoritzLaurer/deberta-v3-base-zeroshot-v2.0"
 file_name = "onnx/model.onnx"
 tokenizer_name = "MoritzLaurer/deberta-v3-base-zeroshot-v2.0"
-model = ORTModelForSequenceClassification.from_pretrained(model_name, file_name=file_name)
+model = ORTModelForSequenceClassification.from_pretrained(model_name, file_name=file_name, provider="CUDAExecutionProvider")
 tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, model_max_length=512)
-classifier = pipeline(task="zero-shot-classification", model=model, tokenizer=tokenizer)
-
-# few_shot_tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-small-en-v1.5', model_max_length=512) # 'BAAI/bge-small-en-v1.5'
-# ort_model = ORTModelForFeatureExtraction.from_pretrained('BAAI/bge-small-en-v1.5', file_name="onnx/model.onnx") # 'BAAI/bge-small-en-v1.5'
-# few_shot_model = SetFitModel.from_pretrained("moshew/bge-small-en-v1.5_setfit-sst2-english") # "moshew/bge-small-en-v1.5_setfit-sst2-english"
-
-# Train few_shot_model
-# candidate_labels = ["supported", "refuted"]
-# reference_dataset = load_dataset("SetFit/sst2")
-# dummy_dataset = Dataset.from_dict({})
-# train_dataset = get_templated_dataset(dummy_dataset, candidate_labels=candidate_labels, sample_size=8, template="The CONCLUSION is {} by the PASSAGE.")
-# args = TrainingArguments(
-#     batch_size=32,
-#     num_epochs=1
-# )
-# trainer = Trainer(
-#     model=few_shot_model,
-#     args=args,
-#     train_dataset=train_dataset,
-#     eval_dataset=reference_dataset["test"]
-# )
-# trainer.train()
-
-# onnx_few_shot_model = OnnxSetFitModel(ort_model, few_shot_tokenizer, few_shot_model.model_head)
-
-
+classifier = pipeline(task="zero-shot-classification", model=model, tokenizer=tokenizer, device="cuda:0")
 
 def classify(data_string, request: gradio.Request):
     if request:

requirements.txt

@@ -2,7 +2,7 @@ fastapi==0.88.0
 huggingface_hub==0.23.5
 json5==0.9.25
 numpy<2.0
-optimum[exporters,onnxruntime]==1.21.4
+optimum[exporters,onnxruntime-gpu]==1.21.4
 setfit==1.0.3
 transformers==4.40.2
 sentence-transformers==3.0.1