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

# + [markdown] id="TA21Jo5d9SVq"
#
#
# ![JohnSnowLabs](https://nlp.johnsnowlabs.com/assets/images/logo.png)
#
# [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/JohnSnowLabs/spark-nlp-workshop/blob/master/tutorials/streamlit_notebooks/T5TRANSFORMER_TRANSLATION.ipynb)
#
#
#

# + [markdown] id="CzIdjHkAW8TB"
# # **Text Translation using google's T5 Transformer**

# + [markdown] id="y1sZkQkt3sXX"
# ### Spark NLP documentation and instructions:
# https://nlp.johnsnowlabs.com/docs/en/quickstart
#
# ### You can find details about Spark NLP annotators here:
# https://nlp.johnsnowlabs.com/docs/en/annotators
#
# ### You can find details about Spark NLP models here:
# https://nlp.johnsnowlabs.com/models
#

# + [markdown] id="wIeCOiJNW-88"
# ## 1. Colab Setup

# + id="CGJktFHdHL1n" colab={"base_uri": "https://localhost:8080/"} outputId="bf39d78d-81c6-4c89-ccb2-3f124e45afda"
# Install java
# ! apt-get update -qq
# ! apt-get install -y openjdk-8-jdk-headless -qq > /dev/null
# ! java -version

# Install pyspark
# ! pip install --ignore-installed -q pyspark==2.4.4

# Install Spark NLP
# ! pip install --ignore-installed spark-nlp==2.7.0

# + [markdown] id="eCIT5VLxS3I1"
# ## 2. Start the Spark session

# + [markdown] id="khjM-z9ORFU3"
# Import dependencies and start Spark session.

# + id="sw-t1zxlHTB7"
import os
import json
os.environ['JAVA_HOME'] = "/usr/lib/jvm/java-8-openjdk-amd64"

from pyspark.ml import Pipeline
from pyspark.sql import SparkSession
import pyspark.sql.functions as F
from sparknlp.annotator import *
from sparknlp.base import *
import sparknlp

spark = sparknlp.start()

# + [markdown] id="9RgiqfX5XDqb"
# ## 3. Select the DL model

# + [markdown] id="TPqBOgDhnvjB"
# For complete model list: 
# https://nlp.johnsnowlabs.com/models
#
# For `T5` models:
# https://nlp.johnsnowlabs.com/models?tag=t5

# + [markdown] id="jQACwlw5dJT6"
# ##4. Text Translation using T5 Transformer - English to German

# + [markdown] id="XftYgju4XOw_"
#  Define Spark NLP pipeline

# + colab={"base_uri": "https://localhost:8080/"} id="lBggF5P8J1gc" outputId="1db841b3-5663-46b5-eeb8-0f87bbc214e4"
from sparknlp.annotator import *
from sparknlp.base import *

from pyspark.ml import Pipeline

document_assembler = DocumentAssembler()\
.setInputCol("text")\
.setOutputCol("documents")

sentence_detector = SentenceDetectorDLModel().pretrained()\
  .setInputCols("documents")\
  .setOutputCol("sentence")
  
t5 = T5Transformer().pretrained("t5_small", 'en') \
  .setInputCols(["sentence"]) \
  .setOutputCol("translation")\
  .setTask("translate English to German:")\
  .setMaxOutputLength(200)
  
pipeline = Pipeline(stages=[
    document_assembler, 
    sentence_detector,
    t5
])

data = spark.createDataFrame([
  [1, "My name is Spark NLP! It's nice to meet you."],
  [2, "My name is Wolfgang and I live in Berlin"]
]).toDF('id', 'text')

results = pipeline.fit(data).transform(data)

results.select("translation.result").show(truncate=False)

# + [markdown] id="MbTZ4OMDZWxe"
# ##5. Text Translation using T5 Transformer - English to French

# + colab={"base_uri": "https://localhost:8080/"} id="9FFPVzI2ZT02" outputId="d869439f-ddd8-471c-bdab-70791a5c6ebc"
from sparknlp.annotator import *
from sparknlp.base import *

from pyspark.ml import Pipeline

document_assembler = DocumentAssembler()\
.setInputCol("text")\
.setOutputCol("documents")

sentence_detector = SentenceDetectorDLModel().pretrained()\
  .setInputCols("documents")\
  .setOutputCol("sentence")
  
t5 = T5Transformer().pretrained("t5_small", 'en') \
  .setInputCols(["sentence"]) \
  .setOutputCol("translation")\
  .setTask("translate English to French:")\
  .setMaxOutputLength(200)
  
pipeline = Pipeline(stages=[
    document_assembler, 
    sentence_detector,
    t5
])

data = spark.createDataFrame([
  [1, "My name is Spark NLP! It's nice to meet you."],
  [2, "My name is Wolfgang and I live in Berlin"]
]).toDF('id', 'text')

results = pipeline.fit(data).transform(data)

results.select("translation.result").show(truncate=False)

StarcoderdataJupyterScriptsDedupFiltered