import gradio as gr
from datasets import load_dataset
from ibm_watsonx_ai import APIClient, Credentials
from ibm_watsonx_ai.foundation_models import ModelInference

def run_eval(ibm_api_key, project_id, subset, task, n_samples):
    try:
        credentials = Credentials(
            url="https://eu-gb.ml.cloud.ibm.com",
            api_key=ibm_api_key
        )
        client = APIClient(credentials)
        
        model = ModelInference(
            model_id="meta-llama/llama-3-3-70b-instruct",
            api_client=client,
            project_id=project_id
        )
        
        ds = load_dataset("RMT-team/babilong", subset)
        samples = ds[task]
        
        results = []
        correct = 0
        n = int(n_samples)
        
        for i in range(n):
            s = samples[i]
            prompt = f"{s['input']}\n\nQuestion: {s['question']}\nAnswer:"
            response = model.generate_text(prompt=prompt, params={"max_new_tokens": 15})
            pred = response.strip().lower()
            gold = s['target'].lower()
            match = gold in pred
            if match:
                correct += 1
            results.append(f"[{i+1}] Gold: {gold} | Pred: {pred} | {'OK' if match else 'WRONG'}")
        
        summary = f"Accuracy: {correct}/{n} = {correct/n:.2%}\n\n" + "\n".join(results)
        return summary
    
    except Exception as e:
        return f"ERROR: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("## BABILong x Granite Evaluator (watsonx)")
    with gr.Row():
        apikey_box = gr.Textbox(label="IBM API Key", type="password")
        projectid_box = gr.Textbox(label="Project ID")
    with gr.Row():
        subset_box = gr.Dropdown(["0k","1k","2k","4k","8k"], label="Context Length", value="0k")
        task_box = gr.Dropdown(["qa1","qa2","qa3","qa4","qa5"], label="Task", value="qa1")
        sample_slider = gr.Slider(5, 50, value=10, step=5, label="Samples")
    run_btn = gr.Button("Run Evaluation")
    output_box = gr.Textbox(label="Results", lines=25)
    run_btn.click(fn=run_eval, inputs=[apikey_box, projectid_box, subset_box, task_box, sample_slider], outputs=output_box)

demo.launch()