from bs4 import BeautifulSoup
import torch, torchaudio, torchtext
from torch import nn


class UltimateMLPipeline(nn.Module):
    """An engineer can dream"""

    def __init__(self, pretrained=True):
        super().__init__()
        self.asr_model = torchaudio.asr.get_sota_model()
        self.ner_model = torchtext.ner.get_sota_model()

    def forward(self, exec_meeting_audio, implementation_url='https://paperswithcode.com'):
        business_needs = self.asr_model(exec_meeting_audio)
        latest_research = self.ner_model(BeautifulSoup(implementation_url, 'html.parser'))
        similarity_score = torch.cdist(latest_research, business_needs.use_case, p=2)
        ml_task = torch.max(similarity_score)
        top_papers = ml_task.filter(language__isin=business_needs.tech_stack).order_by(business_needs.kpi)
        best_model = top_papers[0].implementation.get_model(pretrained=True)
        return best_model(business_needs.input)


model = UltimateMLPipeline()
torch.save(model.state_dict(), 'UltimateMLPipeline.ckpt')  # Weights available tomorrow ;)