SBERT Inference
SBERT inference engine for Frankenstein Transformer.
Provides SBERTInference for computing sentence embeddings,
similarity scores, semantic search, clustering, and embedding
serialization using fine-tuned Sentence-BERT models.
- class src.sbert.inference_sbert.SBERTInference(model_path: str, device: str | None = None, batch_size: int = 32)[source]
Bases:
objectInference engine for fine-tuned SBERT models.
Loads a
SentenceTransformermodel and provides methods for encoding sentences to embeddings, computing pairwise similarity, semantic search over a corpus, clustering, and embedding serialization.- model_path
Path to the fine-tuned SBERT model directory.
- batch_size
Default batch size for encoding.
- device
Resolved PyTorch device string.
- model
Loaded
SentenceTransformerinstance.
- max_seq_length
Maximum sequence length from the loaded model.
- __init__(model_path: str, device: str | None = None, batch_size: int = 32)[source]
Initialize the SBERT inference engine.
- Parameters:
model_path – Path to a fine-tuned SBERT model directory.
device – Device string (
"cuda","cpu", orNonefor auto-resolution).batch_size – Default batch size for encoding operations.
- batch_compare(sentences1: List[str], sentences2: List[str]) List[SimilarityResult][source]
Batch comparison of sentence pairs.
- Parameters:
sentences1 – First sentences
sentences2 – Second sentences (must match length of sentences1)
- Returns:
List of SimilarityResult objects
- cluster_sentences(sentences: List[str], n_clusters: int = 5, method: str = 'kmeans') Tuple[numpy.ndarray, numpy.ndarray][source]
Cluster sentences by semantic similarity.
- Parameters:
sentences – List of sentences to cluster
n_clusters – Number of clusters
method – Clustering method (‘kmeans’ or ‘agglomerative’)
- Returns:
(cluster_labels, embeddings)
- compute_similarity(sentence1: str | List[str], sentence2: str | List[str], metric: str = 'cosine') float | numpy.ndarray[source]
Compute similarity between sentence(s).
- Parameters:
sentence1 – First sentence(s)
sentence2 – Second sentence(s)
metric – Similarity metric (‘cosine’ or ‘dot’)
- Returns:
Similarity score(s) in range [0, 1] (or [-1, 1] for unnormalized)
- encode(sentences: str | List[str], batch_size: int | None = None, show_progress: bool = False, normalize: bool = True) numpy.ndarray[source]
Encode sentences into embeddings.
- Parameters:
sentences – Single sentence or list of sentences
batch_size – Batch size (uses default if None)
show_progress – Show progress bar
normalize – Normalize embeddings to unit length
- Returns:
Embeddings array of shape (n_sentences, embedding_dim)
- find_most_similar(query: str, candidates: List[str], top_k: int = 5) List[Tuple[str, float]][source]
Find most similar sentences to query from candidates.
- Parameters:
query – Query sentence
candidates – List of candidate sentences
top_k – Number of top results to return
- Returns:
List of (sentence, similarity_score) tuples, sorted by similarity
- load_embeddings(input_path: str) Tuple[numpy.ndarray, List[str], dict | None][source]
Load precomputed embeddings.
- Parameters:
input_path – Path to .npz file
- Returns:
(embeddings, sentences, metadata)
- save_embeddings(sentences: List[str], output_path: str, metadata: dict | None = None)[source]
Encode sentences and save embeddings to file.
- Parameters:
sentences – List of sentences
output_path – Output file path (.npz format)
metadata – Optional metadata to save
- semantic_search(queries: str | List[str], corpus: List[str], top_k: int = 5) List[List[Tuple[int, float]]][source]
Semantic search: find most similar corpus sentences for each query.
- Parameters:
queries – Query sentence(s)
corpus – Corpus of sentences to search
top_k – Number of results per query
- Returns:
List of results for each query, each result is (corpus_idx, score)