import gradio as gr
import pandas as pd
import random

# Sample data: Replace this with your legal QA dataset structure
# Assuming columns: 'DocID', 'QueryID', 'Query', 'Segment', 'Label'
sample_data = pd.DataFrame({
    'DocID': ['Doc1', 'Doc2', 'Doc3', 'Doc4', 'Doc5'],
    'QueryID': [101, 102, 103, 104, 105],
    'Query': ['What is the law regarding...', 'How is the case...', 'Definition of legal term...', 'Procedure for filing...', 'Rights of an individual...'],
    'Segment': ['Segment1', 'Segment2', 'Segment3', 'Segment4', 'Segment5'],
    'Label': [1, 0, 1, 0, 1]  # Sample labels
})

# Fake predictions: You should replace these with actual predictions from your test set
fake_predictions = {
    101: 'Positive Response',
    102: 'Negative Response',
    103: 'Positive Response',
    104: 'Negative Response',
    105: 'Positive Response'
}

def predict(query_id):
    # Simulate a model prediction
    response = fake_predictions.get(query_id, "Unknown QueryID")
    return response

def get_random_row():
    # Get a random row from the dataset for demonstration
    random_row = sample_data.sample().iloc[0]
    return f"DocID: {random_row['DocID']}, QueryID: {random_row['QueryID']}, Query: {random_row['Query']}, Segment: {random_row['Segment']}"

iface = gr.Interface(
    fn=predict,
    inputs=gr.inputs.Dropdown(list(sample_data['QueryID']), label="Select QueryID"),
    outputs="text",
    examples=[get_random_row() for _ in range(5)]
)

iface.launch()