import gradio as gr
import pandas as pd
import re
from pyspark.sql import SparkSession, Window
import pyspark.sql.functions as F
from llama_cpp import Llama
from loguru import logger  # Import the logger from loguru

# Create the models directory
!mkdir -p ./models

# Download the Llama model files
!wget -O ./models/llama-2-7b-chat.ggmlv3.q8_0.bin https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGML/resolve/main/llama-2-7b-chat.ggmlv3.q8_0.bin
!wget -O ./models/llama-2-7b-chat.ggmlv3.q2_K.bin https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGML/resolve/main/llama-2-7b-chat.ggmlv3.q2_K.bin

# download "War and Peace" from Project Gutenberg
!mkdir -p ./data
!curl "https://gutenberg.org/cache/epub/2600/pg2600.txt" -o ./data/war_and_peace.txt


# Define the Llama models
MODEL_Q8_0 = Llama(model_path="./models/llama-2-7b-chat.ggmlv3.q8_0.bin", n_ctx=8192, n_batch=512)
MODEL_Q2_K = Llama(model_path="./models/llama-2-7b-chat.ggmlv3.q2_K.bin", n_ctx=8192, n_batch=512)

# Function to read the text file and create Spark DataFrame
def create_spark_dataframe(text):
    # Get list of chapter strings
    chapter_list = [x for x in re.split('CHAPTER .+', text) if len(x) > 100]

    # Create Spark DataFrame
    spark = SparkSession.builder.appName("Counting word occurrences from a book, under a microscope.").config("spark.driver.memory", "4g").getOrCreate()
    spark.sparkContext.setLogLevel("WARN")
    df = spark.createDataFrame(pd.DataFrame({'text': chapter_list, 'chapter': range(1, len(chapter_list) + 1)}))

    return df

# Function to summarize a chapter using the selected model
def llama2_summarize(chapter_text, model_version):
    # Choose the model based on the model_version parameter
    if model_version == "q8_0":
        llm = MODEL_Q8_0
    elif model_version == "q2_K":
        llm = MODEL_Q2_K
    else:
        return "Error: Invalid model_version."

    # Template for this model version
    template = """
    [INST] <<SYS>>
    You are a helpful, respectful and honest assistant. 
    Always answer as helpfully as possible, while being safe.  
    Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. 
    Please ensure that your responses are socially unbiased and positive in nature.
    If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. 
    If you don't know the answer to a question, please don't share false information.
    <</SYS>>
    {INSERT_PROMPT_HERE} [/INST]
    """

    # Create prompt
    prompt = 'Summarize the following novel chapter in a single sentence (less than 100 words): ' + chapter_text
    prompt = template.replace('INSERT_PROMPT_HERE', prompt)

    # Log the input chapter text and model_version
    logger.info(f"Input chapter text: {chapter_text}")
    logger.info(f"Selected model version: {model_version}")

    # Generate summary using the selected model
    output = llm(prompt, max_tokens=-1, echo=False, temperature=0.2, top_p=0.1)
    summary = output['choices'][0]['text']

    # Log the generated summary
    logger.info(f"Generated summary: {summary}")

    return summary

# Read the "War and Peace" text file and create Spark DataFrame
with open('/content/data/war_and_peace.txt', 'r') as file:
    text = file.read()
df_chapters = create_spark_dataframe(text)

# Create summaries via Spark
summaries = (df_chapters
                .limit(1)
                .groupby('chapter')
                .applyInPandas(llama2_summarize, schema='summary string, chapter int')
                .show(vertical=True, truncate=False)
            )

# Prompt for the file
file_path = gr.inputs.File(label="Upload 'War and Peace' text file")

# Choose the model version
model_version = gr.inputs.Radio(["q8_0", "q2_K"], label="Choose Model Version")

# Define the Gradio interface
iface = gr.Interface(
    fn=llama2_summarize,
    inputs=[file_path, model_version],
    outputs="text",  # Summary text
    live=False,
    capture_session=True,
    title="Llama2 Chapter Summarizer",
    description="Upload a text file of the novel 'War and Peace', and choose the model version ('q8_0' or 'q2_K') to get a summarized sentence for each chapter.",
)

iface.launch();