File size: 5,693 Bytes
65c2129 | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | import gradio as gr
from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch
import spacy
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
import subprocess
# Download and load spaCy model
try:
nlp = spacy.load("en_core_web_sm")
except OSError:
print("Downloading spaCy model...")
subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])
nlp = spacy.load("en_core_web_sm")
# Load models
MODEL_NAME = "bert-large-uncased"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForMaskedLM.from_pretrained(MODEL_NAME)
model.eval()
# Optimize for CPU
torch.set_num_threads(4) # Adjust based on available CPU cores
# Load sentence transformer for semantic similarity
similarity_model = SentenceTransformer('all-MiniLM-L6-v2')
# POS tag mapping
POS_MAP = {
'NOUN': ['NOUN', 'PROPN'],
'VERB': ['VERB'],
'ADJ': ['ADJ'],
'ADV': ['ADV'],
'PROPN': ['NOUN', 'PROPN']
}
def get_pos_tag(word):
"""Get POS tag for a word"""
doc = nlp(word)
if len(doc) > 0:
return doc[0].pos_
return None
def suggest_alternatives(text, top_k=5, similarity_threshold=0.3):
if not text.strip():
return "Please enter some text."
words = text.split()
# Get POS tags for all words
doc = nlp(text)
word_pos = {token.text: token.pos_ for token in doc}
# Filter valid words and create masked versions
valid_word_indices = []
masked_texts = []
for i, word in enumerate(words):
if not word.isalpha() or len(word) < 2:
continue
original_pos = word_pos.get(word)
if not original_pos:
continue
# Create masked version
masked_words = words.copy()
masked_words[i] = "[MASK]"
masked_text = " ".join(masked_words)
valid_word_indices.append(i)
masked_texts.append(masked_text)
if not masked_texts:
return "No valid words found to analyze."
# **BATCHED INFERENCE** - Process all masked sentences at once
inputs = tokenizer(masked_texts, return_tensors="pt", padding=True, truncation=True)
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits # Shape: [batch_size, seq_len, vocab_size]
# Process results for each word
suggestions = {}
for batch_idx, word_idx in enumerate(valid_word_indices):
word = words[word_idx]
original_pos = word_pos.get(word)
# Find mask position in this batch item
mask_positions = (inputs["input_ids"][batch_idx] == tokenizer.mask_token_id).nonzero(as_tuple=True)[0]
if len(mask_positions) == 0:
continue
mask_pos = mask_positions[0]
# Get predictions for this masked position
probs = torch.nn.functional.softmax(logits[batch_idx, mask_pos, :], dim=-1)
top_indices = torch.topk(probs, top_k * 10).indices.tolist()
# Decode and filter candidates
candidates = []
for idx in top_indices:
decoded = tokenizer.decode([idx]).strip()
# Basic filters
if not decoded.isalpha() or len(decoded) < 2 or decoded.lower() == word.lower():
continue
# POS filter
candidate_pos = get_pos_tag(decoded)
if candidate_pos:
allowed_pos = POS_MAP.get(original_pos, [original_pos])
if candidate_pos not in allowed_pos:
continue
# Semantic similarity filter
word_embedding = similarity_model.encode([word.lower()])
candidate_embedding = similarity_model.encode([decoded.lower()])
similarity = cosine_similarity(word_embedding, candidate_embedding)[0][0]
if similarity >= similarity_threshold:
candidates.append((decoded, float(similarity)))
if len(candidates) >= top_k:
break
if candidates:
candidates.sort(key=lambda x: x[1], reverse=True)
suggestions[word] = candidates[:top_k]
if not suggestions:
return "No suggestions found. Try lowering the similarity threshold."
# Format output
output = ""
for word, alts in suggestions.items():
pos_tag = word_pos.get(word, 'unknown')
output += f"**{word}** ({pos_tag}) → "
formatted_alts = [f"{alt} `({score:.2f})`" for alt, score in alts]
output += ", ".join(formatted_alts)
output += "\n\n"
return output
demo = gr.Interface(
fn=suggest_alternatives,
inputs=[
gr.Textbox(label="Input text", placeholder="Enter a sentence to analyze"),
gr.Slider(3, 10, value=5, step=1, label="Number of suggestions per word"),
gr.Slider(0.0, 1.0, value=0.3, step=0.05, label="Similarity threshold (higher = more similar)")
],
outputs=gr.Markdown(label="Word Alternatives with Similarity Scores"),
title="BERT Word Alternatives (Optimized)",
description="🚀 **Batched processing** for faster results! Shows contextually relevant word replacements using BERT-Large with POS filtering and semantic similarity scores.",
examples=[
["Artificial intelligence is transforming global commerce.", 5, 0.3],
["The quick brown fox jumps over the lazy dog.", 5, 0.4],
["Machine learning algorithms process data efficiently.", 7, 0.2]
]
)
if __name__ == "__main__":
demo.launch() |