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):
        # expects a list of strings
        if type(txts) == str:
            txts = [txts]
        embedded = self.embedder.encode(np.array(txts))
        embedded_tensor = torch.tensor(embedded, dtype=torch.float32)
        regressed = self.regressor(embedded_tensor)
        val = regressed.flatten().tolist()[0]
        return round(val, 4)