| import torch | |
| import json | |
| from huggingface_hub import hf_hub_download | |
| import re | |
| import emoji | |
| from transformers import BertForSequenceClassification, BertTokenizer | |
| def preprocess_text(text): | |
| """Preprocess the input text to match training conditions.""" | |
| text = re.sub(r'u/\w+', '[USER]', text) | |
| text = re.sub(r'r/\w+', '[SUBREDDIT]', text) | |
| text = re.sub(r'http[s]?://\S+', '[URL]', text) | |
| text = emoji.demojize(text, delimiters=(" ", " ")) | |
| text = text.lower() | |
| return text | |
| def load_model_and_resources(): | |
| """Load the model, tokenizer, emotion labels, and thresholds from Hugging Face.""" | |
| repo_id = "logasanjeev/emotions-analyzer-bert" | |
| try: | |
| model = BertForSequenceClassification.from_pretrained(repo_id) | |
| tokenizer = BertTokenizer.from_pretrained(repo_id) | |
| except Exception as e: | |
| raise RuntimeError(f"Error loading model/tokenizer: {str(e)}") | |
| try: | |
| thresholds_file = hf_hub_download(repo_id=repo_id, filename="optimized_thresholds.json") | |
| with open(thresholds_file, "r") as f: | |
| thresholds_data = json.load(f) | |
| if not (isinstance(thresholds_data, dict) and "emotion_labels" in thresholds_data and "thresholds" in thresholds_data): | |
| raise ValueError("Unexpected format in optimized_thresholds.json. Expected a dictionary with keys 'emotion_labels' and 'thresholds'.") | |
| emotion_labels = thresholds_data["emotion_labels"] | |
| thresholds = thresholds_data["thresholds"] | |
| except Exception as e: | |
| raise RuntimeError(f"Error loading thresholds: {str(e)}") | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| model.to(device) | |
| model.eval() | |
| return model, tokenizer, emotion_labels, thresholds, device | |
| MODEL, TOKENIZER, EMOTION_LABELS, THRESHOLDS, DEVICE = None, None, None, None, None | |
| def predict_emotions(text): | |
| """Predict emotions for the given text using the GoEmotions BERT model. | |
| Args: | |
| text (str): The input text to analyze. | |
| Returns: | |
| tuple: (predictions, processed_text) | |
| - predictions (str): Formatted string of predicted emotions and their confidence scores. | |
| - processed_text (str): The preprocessed input text. | |
| """ | |
| global MODEL, TOKENIZER, EMOTION_LABELS, THRESHOLDS, DEVICE | |
| if MODEL is None: | |
| MODEL, TOKENIZER, EMOTION_LABELS, THRESHOLDS, DEVICE = load_model_and_resources() | |
| processed_text = preprocess_text(text) | |
| encodings = TOKENIZER( | |
| processed_text, | |
| padding='max_length', | |
| truncation=True, | |
| max_length=128, | |
| return_tensors='pt' | |
| ) | |
| input_ids = encodings['input_ids'].to(DEVICE) | |
| attention_mask = encodings['attention_mask'].to(DEVICE) | |
| with torch.no_grad(): | |
| outputs = MODEL(input_ids, attention_mask=attention_mask) | |
| logits = torch.sigmoid(outputs.logits).cpu().numpy()[0] | |
| predictions = [] | |
| for i, (logit, thresh) in enumerate(zip(logits, THRESHOLDS)): | |
| if logit >= thresh: | |
| predictions.append((EMOTION_LABELS[i], round(logit, 4))) | |
| predictions.sort(key=lambda x: x[1], reverse=True) | |
| result = "\n".join([f"{emotion}: {confidence:.4f}" for emotion, confidence in predictions]) or "No emotions predicted." | |
| return result, processed_text | |
| if __name__ == "__main__": | |
| import argparse | |
| parser = argparse.ArgumentParser(description="Predict emotions using the GoEmotions BERT model.") | |
| parser.add_argument("text", type=str, help="The input text to analyze for emotions.") | |
| args = parser.parse_args() | |
| result, processed = predict_emotions(args.text) | |
| print(f"Input: {args.text}") | |
| print(f"Processed: {processed}") | |
| print("Predicted Emotions:") | |
| print(result) |