import gradio as gr
import os
import tensorflow as tf
from transformers import TFAutoModelForQuestionAnswering, AutoTokenizer

# Load text files
def load_text_files_from_directory(directory):
    texts = []
    for filename in os.listdir(directory):
        if filename.endswith(".txt"):
            with open(os.path.join(directory, filename), 'r') as file:
                texts.append(file.read())
    return texts

# Load models and tokenizers
model_options = {
    "DistilRoBERTa": "MuniebAbdelrahman/distilroberta",
    "DistilBERT": "MuniebAbdelrahman/distilbert"
}
models = {}
tokenizers = {}
for name, path in model_options.items():
    models[name] = TFAutoModelForQuestionAnswering.from_pretrained(path)
    tokenizers[name] = AutoTokenizer.from_pretrained(path)

# Answer question function
def answer_question(model_name, question, directory):
    model = models[model_name]
    tokenizer = tokenizers[model_name]
    texts = load_text_files_from_directory(directory)
    
    best_answer = ""
    best_score = -float('inf')
    
    for text in texts:
        inputs = tokenizer.encode_plus(question, text, return_tensors="tf", truncation=True)
        input_ids = inputs["input_ids"].numpy()[0]
        
        outputs = model(inputs)
        answer_start_scores = outputs.start_logits
        answer_end_scores = outputs.end_logits
        
        answer_start = tf.argmax(answer_start_scores, axis=1).numpy()[0]
        answer_end = (tf.argmax(answer_end_scores, axis=1) + 1).numpy()[0]
        
        if answer_start == 0 or answer_end == 0:  # This assumes [CLS] token is at index 0
            continue
        
        score = answer_start_scores[0, answer_start].numpy() + answer_end_scores[0, answer_end].numpy()
        
        if score > best_score:
            best_score = score
            best_answer = tokenizer.decode(input_ids[answer_start:answer_end])
    
    if best_answer == "":
        return "The model couldn't find a suitable answer in the context."
    
    return best_answer

# Gradio interface
def gradio_interface(model_name, question):
    directory = 'data_set'  # Ensure this path matches where you uploaded your text files
    return answer_question(model_name, question, directory)

model_names = list(model_options.keys())
gr.Interface(
    fn=gradio_interface, 
    inputs=[
        gr.Dropdown(choices=model_names, label="Select Model"),
        gr.Textbox(label="Question")
    ], 
    outputs="text"
).launch()