# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Fake Reviews
#
# Data Source
#
# https://www.kaggle.com/yelp-dataset/yelp-dataset
#
# The overall goal of this notebook is to create examples that we can use to train a neural network that will be able to create fake reviews.

#import findspark osx
#findspark.init()
import pyspark

# Create a spark session with Apache Arrow enabled.

from pyspark.sql import SparkSession
spark = SparkSession.builder. \
    appName('big-data') \
    .master("local[8]") \
    .config("spark.driver.memory", "8g") \
    .config('spark.sql.execution.arrow.enabled', 'true').getOrCreate()

# ## Reviews
#
# Loading the yelp reviews, printing their schema and dropping NA rows.

reviews_df = spark.read.json('yelp-dataset/yelp_academic_dataset_review.json')

reviews_df.printSchema()

reviews_df.count()

reviews_df = reviews_df.na.drop()

reviews_df.count()

# ## Businesses

business_df = spark.read.json('yelp-dataset/yelp_academic_dataset_business.json')

business_df.count()

business_df.printSchema()

business_df = business_df.na.drop(subset=['city', 'categories'])

business_df.count()

# We ultimately want to create examples in the following form:
#
# 'RATING CITY CATEGORIES': 'REVIEW TEXT'
#
# If we wanted more granular control over our fake review generation, we could also use something like 'RATING BUSINESS_NAME CITY STATE CATEGORIES'
#
# Let's keep it simple and create a new examples data frame:

reviews = reviews_df.alias('reviews')
business = business_df.alias('business')

# +
from pyspark.sql.functions import concat_ws

examples = reviews.join(business, reviews.business_id == business.business_id) \
    .select(concat_ws(' ', reviews.stars, business.city, business.categories).alias('context'), \
            reviews.text.alias('review'))
    

# -

examples.show()

examples.count()

# ## Tokenization
#
# We need to tokenize our examples so that we can feed them into our neural network. We also need to clean up our examples before tokenizing them. We remove new-lines and non-ascii characters

from pyspark.sql.functions import regexp_replace
examples = examples.withColumn('review', regexp_replace(examples.review, '[\\r\\n]', ' '))

examples = examples.withColumn('review', regexp_replace(examples.review, '[^\x00-\x7F]+', ' '))

# We now want to save the context and review columns as text files.

examples.select(examples.context).write.format('text').save('context.txt')
examples.select(examples.review).write.format('text').save('reviews.txt')

# Making sure that both exports have the same line counts

# !cat context.txt/*.txt | wc -l

# !cat reviews.txt/*.txt | wc -l

# Unfortunately, the default SparkML Tokenizer is rather simple, there is a RegexTokenizer, but crafting hand-written tokenization rules with it is error-prone and cumbersome. We could try to use an industrial-strength Tokenizer from spacy. Implementing an SparkML-Transformer would require a Java counter-part on top of the python implementation, so let's fall back to UDF's instead.
#
# Another alternative is using Apache Arrow.
#
# As it turns out, both ways really slow. The following cells are therefore only used for illustrative purposes. I've tested the UDF approach multiple times and had a lot of time outs or OutOfMemory exceptions.
#
# I therefore had to rely on Stanford's CoreNLP package to tokenize the examples. Please consult the tokenization notebook.

# +
import spacy
from pyspark.sql.functions import udf, pandas_udf
from pyspark.sql.types import ArrayType, StringType

nlp = spacy.load("en_core_web_sm")

@udf(ArrayType(StringType()))
def tokenize(s):
    return [token.text for token in nlp(s)]


# +
import spacy
from pyspark.sql.functions import udf, pandas_udf, PandasUDFType
from pyspark.sql.types import ArrayType, StringType

nlp = spacy.load("en_core_web_sm")

def spacy_tokenize(s):
    return [token.text for token in nlp(s)]

@pandas_udf("string", PandasUDFType.SCALAR)
def tokenize(x):
    return x.apply(spacy_tokenize)


# -

tokenized = examples.select(concat_ws(' ', tokenize(examples.context)).alias('context'), concat_ws(' ', tokenize(examples.review)).alias('review'))

train, val, test = tokenized.randomSplit([0.98, 0.016, 0.004], seed=42)

train.count()
val.count()
test.count()


# +
def write_df(df, file_name):
    df.select(df.context).write.format('text').save(file_name + '_src.txt')
    df.select(df.review).write.format('text').save(file_name + '_tgt.txt')
    

# -

train = train.cache()
val = val.cache()
test = test.cache()

write_df(test, 'test')

write_df(val, 'val')

write_df(train, 'train')

StarcoderdataJupyterScriptsDedupFiltered