# -*- coding: utf-8 -*-
import dataiku
import pandas as pd, numpy as np
from dataiku import pandasutils as pdu
import os
from transformers import pipeline
from models import SentimentAnalyzer


# Read recipe inputs
customer_reviews_prepared = dataiku.Dataset("customer_reviews_prepared")
customer_review_sentiment = dataiku.Dataset("customer_review_sentiment")


analyzer = SentimentAnalyzer(review_col="review_text",column_labels=list(customer_reviews_prepared.get_dataframe(limit=1).columns))


iterrator_df = customer_reviews_prepared.iter_dataframes(chunksize=500)
first_batch = True
mark = 0
with customer_review_sentiment.get_writer() as writer:
    for dataset in iterrator_df:
        dataset["review_score_sentiment"] =  analyzer.predict(dataset.values)
        if first_batch:
            customer_review_sentiment.write_schema_from_dataframe(dataset)
            first_batch = False
        writer.write_dataframe(dataset)
        mark+=1
