Commit 83761a81 authored by Janos Borst's avatar Janos Borst
Browse files

method

parent 673a6c8c
Pipeline #37018 canceled with stages
in 23 seconds
......@@ -18,7 +18,7 @@ class Embedder:
"""
A class for embedding lists of sentences.
"""
def __init__(self, representation, device="cpu", return_device="cpu"):
def __init__(self, representation, method="first_token", device="cpu", return_device="cpu"):
"""
Holding an embedding to embed lists of sentences with huggingface or glove.
......@@ -28,6 +28,7 @@ class Embedder:
ToDo:
- Support Pooling method
- Averaging Method
:param representation: Name of the representation (language model or glove vectors)
:param device: The device to use for the computation of the embeddings
......@@ -39,6 +40,8 @@ class Embedder:
self.device = device
self.emb = self.emb.to(self.device)
self.return_device = return_device
assert method in ("first_token",)
self.method = method
def embed(self, sentences: List, pad = None):
"""
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment