Ramrm/aspect_based_sentiment_analysis

1import torch2import spacy3import subprocess4from transformers import AutoTokenizer, AutoModelForTokenClassification5import gradio as gr6 7# Load spaCy8try:9    nlp = spacy.load("en_core_web_sm")10except:11    subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])12    nlp = spacy.load("en_core_web_sm")13 14# Load model from Hugging Face15MODEL_NAME = "Ramrm/absa-model"16 17tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)18model = AutoModelForTokenClassification.from_pretrained(MODEL_NAME)19 20device = torch.device("cuda" if torch.cuda.is_available() else "cpu")21model.to(device)22 23 24def extract_aspects(predictions):25    aspects = []26    current_aspect = []27    current_sentiment = None28 29    sentiment_map = {30        "POS": "positive",31        "NEG": "negative",32        "NEU": "neutral"33    }34 35    for word, label in predictions:36        if label.startswith("B-"):37            if current_aspect:38                aspects.append({39                    "aspect": " ".join(current_aspect),40                    "sentiment": sentiment_map[current_sentiment]41                })42            current_aspect = [word]43            current_sentiment = label.split("-")[1]44 45        elif label.startswith("I-") and current_aspect:46            current_aspect.append(word)47 48        else:49            if current_aspect:50                aspects.append({51                    "aspect": " ".join(current_aspect),52                    "sentiment": sentiment_map[current_sentiment]53                })54                current_aspect = []55                current_sentiment = None56 57    if current_aspect:58        aspects.append({59            "aspect": " ".join(current_aspect),60            "sentiment": sentiment_map[current_sentiment]61        })62 63    return aspects64 65 66def predict(sentence):67    model.eval()68 69    doc = nlp(sentence)70    tokens = [token.text for token in doc]71 72    inputs = tokenizer(73        tokens,74        is_split_into_words=True,75        return_tensors="pt"76    )77 78    inputs = {k: v.to(device) for k, v in inputs.items()}79 80    with torch.no_grad():81        outputs = model(**inputs)82 83    predictions = outputs.logits.argmax(dim=2)84 85    word_ids = tokenizer(tokens, is_split_into_words=True).word_ids()86 87    final_predictions = []88    previous_word_idx = None89 90    for idx, word_idx in enumerate(word_ids):91        if word_idx is None:92            continue93 94        if word_idx != previous_word_idx:95            label = model.config.id2label[predictions[0][idx].item()]96            final_predictions.append((tokens[word_idx], label))97 98        previous_word_idx = word_idx99 100    aspects = extract_aspects(final_predictions)101 102    # Format output nicely103    if not aspects:104        return "No aspects found"105 106    return "\n".join([f"{a['aspect']} → {a['sentiment']}" for a in aspects])107 108 109# Gradio UI110interface = gr.Interface(111    fn=predict,112    inputs=gr.Textbox(label="Enter sentence"),113    outputs=gr.Textbox(label="Aspect Sentiment"),114    title="Aspect-Based Sentiment Analysis",115    description="Enter a sentence to extract aspects and their sentiment"116)117 118interface.launch()