jskinner215/TAPASxHF2

1from transformers import AutoTokenizer, AutoModel2import torch3 4class EmbeddingGenerator:5    def __init__(self):6        self.model_name = "deepset/all-mpnet-base-v2-table"7        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)8        self.model = AutoModel.from_pretrained(self.model_name)9 10    def generate_embeddings(self, dataframes):11        embeddings = []12        for df in dataframes:13            inputs = self.tokenizer(df.to_string(index=False), return_tensors='pt', truncation=True, padding=True)14            outputs = self.model(**inputs)15            embeddings.append(outputs.last_hidden_state.mean(dim=1).detach().numpy())16        return embeddings17