import gradio as gr
from transformers import AutoModelForQuestionAnswering, AutoTokenizer
import torch
import logging

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

class QuestionAnsweringSystem:
    def __init__(self, model_name):
        """Initialize the QA system with the specified model."""
        logger.info(f"Loading model and tokenizer from {model_name}")
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForQuestionAnswering.from_pretrained(model_name)
        self.max_length = 384
        
    def answer_question(self, context, question):
        """Process the question and context to generate an answer."""
        try:
            # Tokenize input text
            inputs = self.tokenizer(
                question,
                context,
                max_length=self.max_length,
                truncation="only_second",
                padding="max_length",
                return_tensors="pt"
            )
            
            # Generate model predictions
            with torch.no_grad():
                outputs = self.model(**inputs)
            
            # Extract answer span
            answer_start = torch.argmax(outputs.start_logits)
            answer_end = torch.argmax(outputs.end_logits)
            
            # Decode tokens to get the answer
            tokens = inputs.input_ids[0][answer_start:answer_end + 1]
            answer = self.tokenizer.decode(tokens, skip_special_tokens=True)
            
            # Calculate confidence score
            confidence_start = torch.softmax(outputs.start_logits, dim=1).max().item()
            confidence_end = torch.softmax(outputs.end_logits, dim=1).max().item()
            confidence = (confidence_start + confidence_end) / 2
            
            return answer, float(confidence)
            
        except Exception as e:
            logger.error(f"Error in answer generation: {str(e)}")
            return "An error occurred while processing your question.", 0.0

# Initialize the QA system
qa_system = QuestionAnsweringSystem("aman-augurs/bert-fine-tuned-qa3e")  # Replace with your model name

def process_query(context, question):
    """Handle the user query and return formatted results."""
    if not context or not question:
        return "Please provide both a context and a question."
    
    try:
        answer, confidence = qa_system.answer_question(context, question)
        
        if confidence < 0.1:
            return "I'm not confident enough to provide an answer based on the given context."
            
        response = f"Answer: {answer}\nConfidence: {confidence:.2%}"
        return response
        
    except Exception as e:
        logger.error(f"Error processing query: {str(e)}")
        return "An error occurred while processing your request."

# Create the Gradio interface
def create_interface():
    """Create and configure the Gradio interface."""
    return gr.Interface(
        fn=process_query,
        inputs=[
            gr.Textbox(
                label="Context",
                placeholder="Введите контекстный отрывок здесь...",
                lines=10
            ),
            gr.Textbox(
                label="Question",
                placeholder="Введите свой вопрос здесь..."
            )
        ],
        outputs=gr.Textbox(label="Response"),
        title="Вопросно-ответная система",
        description="""Это приложение использует тонко настроенную модель BERT для ответа на вопросы на основе предоставленного контекста. 
и задайте конкретный вопрос об этом.""",
        examples=[
            ["The Golden Gate Bridge is a suspension bridge spanning the Golden Gate strait, the one-mile-wide strait connecting San Francisco Bay and the Pacific Ocean. The structure links the U.S. city of San Francisco, California to Marin County, carrying both U.S. Route 101 and California State Route 1 across the strait.", 
             "How wide is the Golden Gate strait?"],
            ["Python is a high-level, interpreted programming language. Python's design philosophy emphasizes code readability with the use of significant indentation. Its language constructs and object-oriented approach aim to help programmers write clear, logical code for small and large-scale projects.",
             "What is Python's design philosophy?"]
        ],
        theme=gr.themes.Base()
    )

# Create app.py for Hugging Face Spaces
if __name__ == "__main__":
    interface = create_interface()
    interface.launch()