app.py

import os
import gradio as gr
import glob
import json
from LLM_usage import JokeGeneratorBot
from prompt import PROMPT

# Define the base path
base_path = "./Data" # "/mnt/d/Semester7/NLP/RAG/Data"

# Global variables for reuse
bot = None
docs = []

# Function to initialize the bot (called once during deployment)
def initialize_bot(api_key):
    global bot, docs

    # Set the API key
    os.environ['GROQ_API_KEY'] = api_key

    pathes = ["./Data/hate_speech_processed.json", "./Data/reddit_jokes2_processed.json", 
              "./Data/stupidstuff_processed.json", "./Data/wocka_processed.json", 
              "./Data/reddit_jokes1_processed.json"]

    # Load documents (done once)
    if not docs:  # Only load if docs are not already loaded
        for path in pathes:
            with open(path, 'r') as f:
                docs.extend(json.load(f))

    # Initialize the bot
    bot = JokeGeneratorBot(docs, PROMPT)

# Function to handle joke generation
def generate_joke_interface(user_message, regime, bm_koef):
    global bot

    if bm_koef < 0 or bm_koef > 1:
        return "Error: BM25 coefficient should be in the range [0, 1]. Please provide a valid value.", [["", ""]] 

    # Check if the bot is initialized
    if bot is None:
        return "Error: Bot is not initialized. Please provide an API key during deployment.", []

    # Map the selected regime to the corresponding flags
    bm25_only = regime == "BM25 Only"
    semantic_only = regime == "Semantic Only"
    scores_combination = regime == "Scores Combination"

    # Call the bot to generate a joke
    result = bot.generate_joke(
        user_message, n=40, bm25_only=bm25_only, semantic_only=semantic_only, scores_combination=scores_combination, bm_koef=bm_koef
    )

    # Format the context as a DataFrame for better display
    context = []
    for doc in result["Context"]:
        context.extend([[k, v] for k, v in doc.items()])
        context.append(["", ""])  # Add empty space between documents

    return result["Response"], context

# Create a setup interface for API key input
def setup_interface(api_key):
    initialize_bot(api_key)
    return "Joke Generator initialized successfully!"

# Create Gradio interface
setup_demo = gr.Interface(
    fn=setup_interface,
    inputs=[gr.Textbox(label="Enter your GROQ API Key")],
    outputs=[gr.Textbox(label="Setup Status")],
    title="Setup Joke Generator",
    description="Initialize the Joke Generator Bot by providing the GROQ API key. \
        (If there is a connection error(on this or next tab) reload the page, wait 5-10 imnutes, \
        reload the page again and reinitialize the joke generator with the API KEY)\
        If you see some runtime error like memory limit exceeded, tell me on mail: vasyarusynb@gmail.com(I can see your email not so fast)\
        or tg: @Beav3rrr and I will redeploy or turn on new instance",
)

regime_options = ["BM25 Only", "Semantic Only", "Scores Combination"]
joke_demo = gr.Interface(
    fn=generate_joke_interface,
    inputs=[
        gr.Textbox(label="Enter your message"),
        gr.Radio(choices=regime_options, label="Choose Regime", value="Scores Combination"),
        gr.Slider(minimum=0, maximum=1, step=0.01, value=0.75, label="BM25 Coefficient"),
    ],
    outputs=[
        gr.Textbox(label="Generated Joke"),
        gr.Dataframe(headers=["Key", "Value"], label="Context"),
    ],
    title="Joke Generator",
    description="Generate jokes based on your input message(Only in English :( )). Select a retrieval regime and view the context used.\
                Be careful, the jokes can be offensive! Try to write a message that is related to the joke you want to hear.\
                (tell me a joke and its title about... or tell me a joke and its title about... it should be a oneliner, dark, pervy, etc.). Sometimes bot works bad :(\
                In this case, try to rewrite a message and send again.\
                Or try to change the regime or BM25 Coefficient.\
                BM25 Coefficient is used to balance the BM25 and semantic scores(It is active only in Scores Combination mode). Semantic scores are multiplied by (1 - BM25 Coefficient).\
                If you want to use only BM25 or semantic scores, select the corresponding regime or set it to 0.0 or 1.0. respectively.",
)

# In order to divide a big string line into a couple of lines, I can place a backslash at the end of the line 

# Combine setup and main interfaces into a tabbed app
demo = gr.TabbedInterface(
    [setup_demo, joke_demo],
    ["Setup", "Joke Generator"]
)

# Launch the interface
demo.launch()
# demo.launch(share=True)