Skip to content

Commit

Permalink
reformat utils.py
Browse files Browse the repository at this point in the history
  • Loading branch information
derixu committed Dec 26, 2024
1 parent d45325c commit b3676a6
Showing 1 changed file with 5 additions and 4 deletions.
9 changes: 5 additions & 4 deletions fastchat/serve/monitor/classify/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,15 +66,16 @@ def chat_completion_openai(model, messages, temperature, max_tokens, api_dict=No
class HuggingFaceClassifier:
def __init__(self, model_path, device=None):
print("Loading model and tokenizer...")
self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.device = device or torch.device(
"cuda" if torch.cuda.is_available() else "cpu"
)
self.pipeline = pipeline(
"text-classification",
model=model_path,
tokenizer=model_path,
device=self.device
device=self.device,
)

def classify_batch(self, input_texts, batch_size=8):
results = self.pipeline(input_texts, batch_size=batch_size, truncation=True)
return [res['label'] == "LABEL_1" for res in results]

return [res["label"] == "LABEL_1" for res in results]

0 comments on commit b3676a6

Please sign in to comment.