import gradio as gr
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
import torch
import logging
import sys

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Add a handler to write logs to a file
file_handler = logging.FileHandler('app.log')
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
logger.addHandler(file_handler)

logger.info("Starting the application")

try:
    logger.info("Loading the dataset")
    ds = load_dataset("knowrohit07/gita_dataset")
    logger.info("Dataset loaded successfully")
except Exception as e:
    logger.error(f"Error loading dataset: {str(e)}")
    sys.exit(1)

try:
    logger.info("Loading the model and tokenizer")
    model_name = "deepset/roberta-base-squad2"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForQuestionAnswering.from_pretrained(model_name)
    logger.info("Model and tokenizer loaded successfully")
except Exception as e:
    logger.error(f"Error loading model or tokenizer: {str(e)}")
    sys.exit(1)

def clean_answer(answer):
    # Remove special tokens and leading/trailing whitespace
    special_tokens = set(tokenizer.all_special_tokens)
    cleaned_answer = ' '.join(token for token in answer.split() if token not in special_tokens)
    return cleaned_answer.strip()

def answer_question(question, system_prompt, temperature, max_new_tokens, top_p, frequency_penalty, presence_penalty, top_k, echo, best_of):
    logger.info(f"Received question: {question}")
    logger.info(f"Parameters: temp={temperature}, max_tokens={max_new_tokens}, top_p={top_p}, freq_penalty={frequency_penalty}, pres_penalty={presence_penalty}, top_k={top_k}, echo={echo}, best_of={best_of}")

    try:
        logger.info("Combining text from dataset")
        context = " ".join([item['Text'] for item in ds['train']])
        logger.info(f"Combined context length: {len(context)} characters")

        logger.info("Tokenizing input")
        inputs = tokenizer.encode_plus(question, context, return_tensors="pt", max_length=512, truncation=True)
        logger.info(f"Input tokens shape: {inputs['input_ids'].shape}")

        logger.info("Getting model output")
        outputs = model(**inputs)
        logger.info(f"Output logits shapes: start={outputs.start_logits.shape}, end={outputs.end_logits.shape}")

        logger.info("Processing output to get answer")
        answer_start = torch.argmax(outputs.start_logits)
        answer_end = torch.argmax(outputs.end_logits) + 1
        raw_answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end]))
        answer = clean_answer(raw_answer)
        logger.info(f"Generated answer: {answer}")

        if not answer:
            logger.warning("Generated answer was empty after cleaning")
            answer = "I'm sorry, but I couldn't find a specific answer to that question based on the Bhagavad Gita. Could you please rephrase your question or ask about one of the core concepts like dharma, karma, bhakti, or the different types of yoga discussed in the Gita?"

        disclaimer = "\n\nPlease note: This response is generated by an AI model based on the Bhagavad Gita. For authoritative information, please consult the original text or scholarly sources."
        full_response = answer + disclaimer
        logger.info("Answer generated successfully")

        return full_response

    except Exception as e:
        logger.error(f"Error in answer_question function: {str(e)}")
        return "I'm sorry, but an error occurred while processing your question. Please try again later."

logger.info("Setting up Gradio interface")
iface = gr.Interface(
    fn=answer_question,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter your question here..."),
        gr.Textbox(lines=2, placeholder="System prompt (optional)"),
        gr.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label="Temperature"),
        gr.Slider(minimum=1, maximum=500, step=1, value=250, label="Max new tokens"),
        gr.Slider(minimum=0, maximum=1, step=0.05, value=0.95, label="Top p"),
        gr.Slider(minimum=0, maximum=2, step=0.1, value=1, label="Frequency penalty"),
        gr.Slider(minimum=0, maximum=2, step=0.1, value=1, label="Presence penalty"),
        gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Top k"),
        gr.Checkbox(label="Echo"),
        gr.Slider(minimum=1, maximum=5, step=1, value=1, label="Best of")
    ],
    outputs="text",
    title="Bhagavad Gita Q&A",
    description="Ask a question about the Bhagavad Gita, and get an answer based on the dataset."
)

logger.info("Launching the Gradio app")
iface.launch()
logger.info("Gradio app launched successfully")