from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
from datasets import load_dataset

# Model details
MODEL_NAME = "Pisethan/sangapac-math"

# Load model and tokenizer
try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
    classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)
except Exception as e:
    classifier = None
    print(f"Error loading model or tokenizer: {e}")

# Load dataset dynamically from Hugging Face or locally
try:
    dataset = load_dataset("Pisethan/sangapac-math-dataset")["train"]  # Load your dataset
    dataset_dict = {entry["input"]: entry for entry in dataset}  # Create a dictionary for lookup
except Exception as e:
    dataset_dict = {}
    print(f"Error loading dataset: {e}")


def predict(input_text):
    if classifier is None:
        return "Model not loaded properly.", {"Error": "Model not loaded properly."}

    try:
        # Predict the category
        result = classifier(input_text)
        label = result[0]["label"]
        score = result[0]["score"]

        # Retrieve output and metadata dynamically from the dataset
        data = dataset_dict.get(input_text, {"output": "Unknown", "metadata": {}})
        output = data["output"]
        metadata = data["metadata"]

        # Create a simple result string
        simple_result = f"Category: {label}\nConfidence: {score:.2f}\nResult: {output}"

        # Create the full JSON output
        detailed_result = {
            "Category": label,
            "Confidence": score,
            "Output (Result)": output,
            "Metadata": metadata,
        }

        return simple_result, detailed_result
    except Exception as e:
        return "An error occurred.", {"Error": str(e)}

# Gradio interface
import gradio as gr

interface = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=2, placeholder="Enter a math problem..."),
    outputs=[
        gr.Textbox(label="Simple Output"),  # Display only the result
        gr.JSON(label="Detailed JSON Output"),  # Display full JSON
    ],
    title="Sangapac Math Model",
    description="A model to classify math problems into categories like Arithmetic, Multiplication, Division, Algebra, and Geometry.",
)

# Launch the app
interface.launch()