Spaces:
Running
Running
import torch | |
import os | |
from sentence_transformers import CrossEncoder | |
from gliclass import GLiClassModel, ZeroShotClassificationPipeline | |
from transformers import AutoTokenizer | |
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') | |
model = GLiClassModel.from_pretrained(os.getenv("GLICLASS_MODEL_PATH")).eval().to(device) | |
tokenizer = AutoTokenizer.from_pretrained(os.getenv("GLICLASS_MODEL_PATH")) | |
multi_label_pipeline = ZeroShotClassificationPipeline(model, tokenizer, classification_type='multi-label', | |
device=device) | |
st = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2") | |