Spaces:
Sleeping
Sleeping
| from sentence_transformers import SentenceTransformer | |
| def load_model(model_name: str, device: str) -> SentenceTransformer: | |
| """Load a SentenceTransformer model onto the specified device. | |
| Args: | |
| model_name: HuggingFace model ID (e.g. "sentence-transformers/all-MiniLM-L6-v2") | |
| device: torch device string ("cpu", "mps", "cuda") | |
| Returns: | |
| Loaded SentenceTransformer model ready for inference | |
| """ | |
| model = SentenceTransformer(model_name, device=device) | |
| return model | |
| def predict(model: SentenceTransformer, texts: list[str]) -> list[list[float]]: | |
| """Generate embeddings for a list of text strings. | |
| Args: | |
| model: Loaded SentenceTransformer model | |
| texts: List of strings to embed | |
| Returns: | |
| List of embedding vectors, each a list of floats | |
| """ | |
| embeddings = model.encode(texts).tolist() | |
| return embeddings | |