Spaces:
Sleeping
Sleeping
File size: 1,699 Bytes
7731b47 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
def classify_with_transformer(text, task="sentiment", model_name="distilbert-base-uncased"):
"""
Classify text using a pre-trained transformer model (BERT, RoBERTa, etc.)
Args:
text (str): Text to analyze
task (str): Classification task ('sentiment', 'emotion', etc.)
model_name (str): Name of the pre-trained model to use
Returns:
dict: Classification results with labels and scores
"""
try:
from transformers import pipeline
# Map tasks to appropriate models if not specified
task_model_map = {
"sentiment": "distilbert-base-uncased-finetuned-sst-2-english",
"emotion": "j-hartmann/emotion-english-distilroberta-base",
"toxicity": "unitary/toxic-bert"
}
# Use mapped model if using default and task is in the map
if model_name == "distilbert-base-uncased" and task in task_model_map:
model_to_use = task_model_map[task]
else:
model_to_use = model_name
# Initialize the classification pipeline
classifier = pipeline(task, model=model_to_use)
# Get classification results
results = classifier(text)
# Format results based on return type (list or dict)
if isinstance(results, list):
if len(results) == 1:
return results[0]
return results
return results
except ImportError:
return {"error": "Required packages not installed. Please install transformers and torch."}
except Exception as e:
return {"error": f"Classification failed: {str(e)}"} |