nextus_regressor_class.py

import torch
from torch import nn
from sentence_transformers import SentenceTransformer
from regressor import *
import numpy as np
import os

ENCODER = os.getenv("ENCODER")

class NextUsRegressor(nn.Module):
    def __init__(self):
        super(NextUsRegressor, self).__init__()
        self.embedder = SentenceTransformer(ENCODER)
        
        
        self.regressor = WRegressor()

        return
    
    def forward(self, txts):
        if type(txts) == str:
            txts = [txts]
        embedded = self.embedder.encode(np.array(txts))
        # embedded_tensor = self.embedder(np.array(txts))
        
        embedded_tensor = torch.tensor(embedded, dtype=torch.float32)
        regressed = self.regressor(embedded_tensor)
        
        # return regressed.tolist()
        # TODO: actually handle list of strings
        vals = regressed.flatten().tolist()
        # must return the whole thing, not just the 0-th element
        strs = list()
        for t, v in list(zip(txts, vals)):
            strs.append(str(round(v, 4)) + "\t" + t[:20])
        return "\n".join(strs)
        # return torch.tensor(val).unsqueeze(1)