import gradio as gr
from sparknlp.base import *
from sparknlp.annotator import *
from sparknlp.pretrained import PretrainedPipeline
import sparknlp

print("Sparknlp Version: " + sparknlp.version())

# Start SparkSession with Spark NLP
# start() functions has 4 parameters: gpu, spark23, spark24, and memory
# sparknlp.start(gpu=True) will start the session with GPU support
# sparknlp.start(spark23=True) is when you have Apache Spark 2.3.x installed
# sparknlp.start(spark24=True) is when you have Apache Spark 2.4.x installed
# sparknlp.start(memory="16G") to change the default driver memory in SparkSession
spark = sparknlp.start()

# Download a pre-trained pipeline
pipeline = PretrainedPipeline('explain_document_dl', lang='en')

# Your testing dataset
placeholder = """
The Mona Lisa is a 16th century oil painting created by Leonardo.
It's held at the Louvre in Paris.
"""

print(spark)

def fn(text: str):
  result = pipeline.annotate(text)
  return result

iface = gr.Interface(
  fn=fn,
  inputs="text",
  outputs="json",
  title="Spark NLP explain_document_dl pipeline",
  description=f"Spark object: {spark}",
  examples=[placeholder],
)

if __name__ == "__main__":
    iface.launch()