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

# 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

# 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)

# 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 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

# Define the Gradio interface
iface = gr.Interface(
    fn=llama2_summarize,
    inputs=[
        gr.inputs.File(label="Upload Text File"),
        "text",
    ],  # chapter_text, model_version
    outputs="text",  # Summary text
    live=False,
    capture_session=True,
    title="Llama2 Chapter Summarizer",
    description="Upload the text file or enter the chapter text and model version ('q8_0' or 'q2_K') to get a summarized sentence.",
)

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