import gradio as gr
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
import re
import logging

logging.basicConfig(level=logging.INFO)

def initialize_pipeline(token):
    """Initialize the Hugging Face pipeline with authentication token."""
    try:
        tokenizer = AutoTokenizer.from_pretrained("riotu-lab/ArabianGPT-0.8B-IslamicQA", use_auth_token=token)
        model = AutoModelForCausalLM.from_pretrained("riotu-lab/ArabianGPT-0.8B-IslamicQA", use_auth_token=token)
        return pipeline("text-generation", model=model, tokenizer=tokenizer)
    except Exception as e:
        logging.error(f"Failed to load model with token. Error: {str(e)}")
        raise

def generate_response(token, question, penalty_alpha, do_sample, top_k, temperature, repetition_penalty, max_new_tokens):
    """Generate the response using the configured parameters and provided question."""
    try:
        # Initialize the pipeline with the token
        pipe = initialize_pipeline(token)
        
        # Apply the formatted prompt
        formatted_prompt_text = f"user\n{question}\nassistant:"
        
        # Generate the response using the model pipeline with specified configuration
        result = pipe(formatted_prompt_text, 
                      num_return_sequences=1,
                      do_sample=do_sample,
                      top_k=top_k,
                      temperature=temperature,
                      repetition_penalty=repetition_penalty,
                      max_new_tokens=max_new_tokens)

        # Assuming the result is in the correct format, extract the generated text
        response_text = result[0]['generated_text'] if result else ""
        
        # Extract and clean the response
        cleaned_response = extract_and_clean_response(response_text)

        return cleaned_response
    except Exception as e:
        logging.error(f"An error occurred during generation: {str(e)}")
        return "An error occurred. Please check your input and try again."

def extract_and_clean_response(text: str) -> str:
    """Extracts the part after 'assistant:' and cleans unwanted patterns from it."""
    if 'assistant:' in text:
        text = text.split('assistant:', 1)[1]
    patterns_to_remove = [
        r'>\s*assistant\s*',  # Remove '> assistant' and similar patterns
        r'<\|[\w\s]+\|>'      # Remove special tokens like '<|...|>'
    ]
    for pattern in patterns_to_remove:
        text = re.sub(pattern, '', text)
    text = text.strip()
    return text

# Setup the Gradio interface
iface = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(placeholder="Enter your Hugging Face token here...", label="Hugging Face Token", type="password"),
        gr.Textbox(lines=2, placeholder="Enter your question here..."),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label="Penalty Alpha"),
        gr.Checkbox(label="Do Sample"),
        gr.Slider(minimum=1, maximum=100, step=1, label="Top K"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label="Temperature"),
        gr.Slider(minimum=1.0, maximum=5.0, step=0.1, label="Repetition Penalty"),
        gr.Slider(minimum=1, maximum=512, step=1, label="Max New Tokens")
    ],
    outputs="text",
    title="ArabianGPT Islamic QA Playground",
    description="Explore the capabilities of ArabianGPT models. Adjust the hyperparameters to see how they affect QA Results",
    live=False  # Requires user to press "submit" to run
)

# Launch the app
iface.launch()