Spaces:

kalyan-ks
/

Ammu_1.0_FactCheckNeed_Classification

Running on Zero

Ammu_1.0_FactCheckNeed_Classification / app.py

Update app.py

b537990 verified 9 days ago

12.2 kB

"""
Ammu-1.0 FactCheck Classifier - Gradio Demo for Hugging Face Spaces
Author: Kalyan KS (https://huggingface.co/kalyan-ks)
"""

import os
import pandas as pd
import gradio as gr
import torch
from transformers import pipeline

# Support Hugging Face ZeroGPU Spaces with local fallback
try:
    import spaces
except ImportError:
    class spaces:
        @staticmethod
        def GPU(func):
            return func

############################################################
# Model Catalog & Pipeline Cache
############################################################

MODEL_OPTIONS = {
    "Ammu-1.0-Base (150M)": "kalyan-ks/Ammu-1.0-FactCheck-Need-base",
    "Ammu-1.0-Small (68M)": "kalyan-ks/Ammu-1.0-FactCheck-Need-small",
    "Ammu-1.0-Tiny (33M)": "kalyan-ks/Ammu-1.0-FactCheck-Need-tiny",
    "Ammu-1.0-Nano (17M)": "kalyan-ks/Ammu-1.0-FactCheck-Need-nano",
}

MODEL_URLS = {
    "Ammu-1.0-Base (150M)": "https://huggingface.co/kalyan-ks/Ammu-1.0-FactCheck-Need-base",
    "Ammu-1.0-Small (68M)": "https://huggingface.co/kalyan-ks/Ammu-1.0-FactCheck-Need-small",
    "Ammu-1.0-Tiny (33M)": "https://huggingface.co/kalyan-ks/Ammu-1.0-FactCheck-Need-tiny",
    "Ammu-1.0-Nano (17M)": "https://huggingface.co/kalyan-ks/Ammu-1.0-FactCheck-Need-nano",
}

MODEL_PERFORMANCE = {
    "Ammu-1.0-Base (150M)": "96.56% Macro-F1",
    "Ammu-1.0-Small (68M)": "96.13% Macro-F1",
    "Ammu-1.0-Tiny (33M)": "96.35% Macro-F1",
    "Ammu-1.0-Nano (17M)": "95.66% Macro-F1",
}

# In-memory pipeline cache to prevent reloading overhead
MODELS = {}

def get_factcheck_pipeline(model_name):
    if model_name not in MODELS:
        device = 0 if torch.cuda.is_available() else -1
        MODELS[model_name] = pipeline(
            task="text-classification",
            model=model_name,
            device=device,
            top_k=None  # Returns probabilities for all classes
        )
    return MODELS[model_name]


############################################################
# Benchmark Examples
############################################################

EXAMPLES = {
    "Example 1 (Factual Question)": "India is the most populous country in the world. Do you agree?",
    "Example 2 (Complex Age Calculation)": "Assume that I was born 16 years 4 months 3 days exactly after Sachin Tendulkar was born. How old I am today?",
    "Example 3 (Short Age Query)": "I was born 16 years 4 months 3 days exactly after Sachin Tendulkar was born. My age?",
    "Example 4 (Input-Text Entity Extraction)": "Kalyan KS is from India. Extract named entities from the given sentence.",
    "Example 5 (Task Instruction with Input)": "Extract named entities: Kalyan KS is from India.",
}


############################################################
# UI & Helper Functions
############################################################

def update_model_info(selected_model_key):
    url = MODEL_URLS[selected_model_key]
    model_id = MODEL_OPTIONS[selected_model_key]
    perf = MODEL_PERFORMANCE[selected_model_key]
    
    info = f"""
- **Model:** [{model_id}]({url})
- **Architecture:** ModernBERT Encoder
- **Parameters:** {selected_model_key.split('(')[-1].replace(')', '')}
- **Task:** Fact-Check Need Classification (Single-Label)
- **Classes:** `FACT_CHECK_NEEDED`, `NO_FACT_CHECK_NEEDED`
- **Performance:** {perf}
"""
    footer = f"""
### 🔗 Resources & Links
- **Hugging Face Model:** [{model_id}]({url})
- **Creator:** [Kalyan KS](https://huggingface.co/kalyan-ks)
- **License:** Apache 2.0
---
Built with ❤️ using **Gradio** and **Hugging Face Transformers**.
"""
    return info, footer


def update_stats(text):
    if not text:
        return "Characters: 0", "Approx. Tokens: 0"
    char_count = len(text)
    approx_tokens = int(len(text.split()) * 1.3)
    return f"Characters: {char_count}", f"Approx. Tokens: {approx_tokens}"


def load_example_text(example_key):
    return EXAMPLES.get(example_key, "")


def clear_inputs():
    empty_df = pd.DataFrame(columns=["Class Label", "Probability", "Confidence Bar"])
    return (
        "",                          # input_text
        "Characters: 0",             # char_caption
        "Approx. Tokens: 0",         # token_caption
        "",                          # status_output
        "",                          # result_card_html
        empty_df,                    # prob_table
        "",                          # metric_prediction
        "",                          # metric_confidence
        "",                          # metric_routing
    )


@spaces.GPU
def process_classification(input_text, selected_model_key):
    if not input_text or not input_text.strip():
        empty_df = pd.DataFrame(columns=["Class Label", "Probability", "Confidence Bar"])
        return (
            "⚠️ Please enter a prompt to classify.",
            "",
            empty_df,
            "-",
            "-",
            "-",
        )

    model_id = MODEL_OPTIONS[selected_model_key]
    classifier = get_factcheck_pipeline(model_id)

    # Run inference
    raw_results = classifier(input_text.strip())

    if isinstance(raw_results, list) and len(raw_results) > 0 and isinstance(raw_results[0], list):
        scores_list = raw_results[0]
    elif isinstance(raw_results, list):
        scores_list = raw_results
    else:
        scores_list = [raw_results]

    scores_dict = {item["label"]: float(item["score"]) for item in scores_list}
    
    top_label = max(scores_dict, key=scores_dict.get)
    top_score = scores_dict[top_label]

    # Styling based on prediction
    if top_label == "FACT_CHECK_NEEDED":
        badge_bg = "#FFEBEE"
        badge_border = "#EF5350"
        badge_color = "#C62828"
        icon = "⚠️"
        action_title = "FACT CHECK NEEDED"
        routing_desc = "Route to Retrieval (RAG) / Search / Knowledge Base before generating response."
        routing_short = "Trigger RAG / Web Search"
    else:
        badge_bg = "#E8F5E9"
        badge_border = "#66BB6A"
        badge_color = "#2E7D32"
        icon = "✅"
        action_title = "NO FACT CHECK NEEDED"
        routing_desc = "Safe for direct LLM generation (Creative writing, code, math logic, or input-grounded task)."
        routing_short = "Direct LLM Generation"

    # HTML Result Card
    styled_html = f"""
    <div style="
        border: 2px solid {badge_border};
        background: {badge_bg};
        border-radius: 12px;
        padding: 20px;
        margin-top: 10px;
        box-shadow: 0 2px 8px rgba(0,0,0,0.05);">
        <div style="display: flex; align-items: center; justify-content: space-between; flex-wrap: wrap; gap: 10px;">
            <div style="display: flex; align-items: center; gap: 12px;">
                <span style="font-size: 28px;">{icon}</span>
                <div>
                    <h3 style="margin: 0; color: {badge_color}; font-size: 20px; font-weight: 700; letter-spacing: 0.5px;">
                        {action_title}
                    </h3>
                    <p style="margin: 4px 0 0 0; color: #424242; font-size: 14px;">
                        {routing_desc}
                    </p>
                </div>
            </div>
            <div style="text-align: right; min-width: 120px;">
                <span style="font-size: 26px; font-weight: 800; color: {badge_color};">
                    {top_score * 100:.2f}%
                </span>
                <div style="font-size: 11px; text-transform: uppercase; color: #616161; font-weight: 600;">
                    Confidence
                </div>
            </div>
        </div>
    </div>
    """

    # Probability Table
    prob_rows = []
    for label, score in sorted(scores_dict.items(), key=lambda x: x[1], reverse=True):
        bar_len = int(score * 20)
        prob_rows.append({
            "Class Label": label,
            "Probability (%)": f"{score * 100:.2f}%",
            "Confidence Bar": "█" * bar_len + "░" * (20 - bar_len),
        })
    prob_df = pd.DataFrame(prob_rows)

    status_msg = f"✅ Classification completed using **{selected_model_key}**."

    return (
        status_msg,
        styled_html,
        prob_df,
        top_label,
        f"{top_score * 100:.2f}%",
        routing_short,
    )


############################################################
# Gradio Interface
############################################################

css = """
#root, body, html, .gradio-container {
    overflow-y: auto !important;
    height: 100vh !important;
}
textarea { font-size: 15px !important; }
"""

with gr.Blocks(title="🔎 Fact-Check Classifier Demo", css=css) as demo:

    gr.Markdown("# 🔎 Fact-Check Need Classification using Ammu-1.0")
    gr.Markdown(
        "Determine whether a prompt requires factual verification (**RAG / Web Search**) or can be processed directly "
        "(**Creative Writing, Reasoning, Code, or Input-Grounded Tasks**) using the fine-tuned **Ammu-1.0 FactCheck** encoder models."
    )

    with gr.Sidebar():
        gr.Markdown("## ⚙️ Settings")
        selected_model = gr.Dropdown(
            choices=list(MODEL_OPTIONS.keys()),
            value="Ammu-1.0-Base (150M)",
            label="Select Model"
        )

        gr.Markdown("---")
        gr.Markdown("### 📌 Model Information")
        model_info_md = gr.Markdown()

    with gr.Column():
        gr.Markdown("### 📝 Input Prompt")

        with gr.Row():
            example_selector = gr.Dropdown(
                choices=list(EXAMPLES.keys()),
                value="Example 1 (Factual Question)",
                label="Choose Example"
            )
            load_example_btn = gr.Button("📄 Load Example")

        input_text = gr.Textbox(
            value=EXAMPLES["Example 1 (Factual Question)"],
            lines=5,
            placeholder="Type or paste a user prompt to classify...",
            label="Enter Prompt"
        )

        with gr.Row():
            char_caption = gr.Markdown("Characters: 0")
            token_caption = gr.Markdown("Approx. Tokens: 0")

        with gr.Row():
            classify_btn = gr.Button("🔍 Classify Prompt", variant="primary")
            clear_btn = gr.Button("🗑️ Clear")

        status_output = gr.Markdown()

    with gr.Column():
        gr.Markdown("---")
        gr.Markdown("### 🎯 Classification Result")
        result_card_html = gr.HTML()

        with gr.Row():
            metric_prediction = gr.Textbox(label="Predicted Class", interactive=False)
            metric_confidence = gr.Textbox(label="Confidence", interactive=False)
            metric_routing = gr.Textbox(label="Routing Recommendation", interactive=False)

        gr.Markdown("### 📊 Class Probability Breakdown")
        prob_table = gr.Dataframe(interactive=False)

    gr.Markdown("---")
    footer_md = gr.Markdown()

    # Event Connections
    demo.load(
        fn=update_model_info,
        inputs=[selected_model],
        outputs=[model_info_md, footer_md]
    )

    demo.load(
        fn=update_stats,
        inputs=[input_text],
        outputs=[char_caption, token_caption]
    )

    selected_model.change(
        fn=update_model_info,
        inputs=[selected_model],
        outputs=[model_info_md, footer_md]
    )

    input_text.change(
        fn=update_stats,
        inputs=[input_text],
        outputs=[char_caption, token_caption]
    )

    load_example_btn.click(
        fn=load_example_text,
        inputs=[example_selector],
        outputs=[input_text]
    )

    clear_btn.click(
        fn=clear_inputs,
        outputs=[
            input_text,
            char_caption,
            token_caption,
            status_output,
            result_card_html,
            prob_table,
            metric_prediction,
            metric_confidence,
            metric_routing,
        ]
    )

    classify_btn.click(
        fn=process_classification,
        inputs=[input_text, selected_model],
        outputs=[
            status_output,
            result_card_html,
            prob_table,
            metric_prediction,
            metric_confidence,
            metric_routing,
        ]
    )

if __name__ == "__main__":
    demo.launch()

	"""
	Ammu-1.0 FactCheck Classifier - Gradio Demo for Hugging Face Spaces
	Author: Kalyan KS (https://huggingface.co/kalyan-ks)
	"""

	import os
	import pandas as pd
	import gradio as gr
	import torch
	from transformers import pipeline

	# Support Hugging Face ZeroGPU Spaces with local fallback
	try:
	import spaces
	except ImportError:
	class spaces:
	@staticmethod
	def GPU(func):
	return func

	############################################################
	# Model Catalog & Pipeline Cache
	############################################################

	MODEL_OPTIONS = {
	"Ammu-1.0-Base (150M)": "kalyan-ks/Ammu-1.0-FactCheck-Need-base",
	"Ammu-1.0-Small (68M)": "kalyan-ks/Ammu-1.0-FactCheck-Need-small",
	"Ammu-1.0-Tiny (33M)": "kalyan-ks/Ammu-1.0-FactCheck-Need-tiny",
	"Ammu-1.0-Nano (17M)": "kalyan-ks/Ammu-1.0-FactCheck-Need-nano",
	}

	MODEL_URLS = {
	"Ammu-1.0-Base (150M)": "https://huggingface.co/kalyan-ks/Ammu-1.0-FactCheck-Need-base",
	"Ammu-1.0-Small (68M)": "https://huggingface.co/kalyan-ks/Ammu-1.0-FactCheck-Need-small",
	"Ammu-1.0-Tiny (33M)": "https://huggingface.co/kalyan-ks/Ammu-1.0-FactCheck-Need-tiny",
	"Ammu-1.0-Nano (17M)": "https://huggingface.co/kalyan-ks/Ammu-1.0-FactCheck-Need-nano",
	}

	MODEL_PERFORMANCE = {
	"Ammu-1.0-Base (150M)": "96.56% Macro-F1",
	"Ammu-1.0-Small (68M)": "96.13% Macro-F1",
	"Ammu-1.0-Tiny (33M)": "96.35% Macro-F1",
	"Ammu-1.0-Nano (17M)": "95.66% Macro-F1",
	}

	# In-memory pipeline cache to prevent reloading overhead
	MODELS = {}

	def get_factcheck_pipeline(model_name):
	if model_name not in MODELS:
	device = 0 if torch.cuda.is_available() else -1
	MODELS[model_name] = pipeline(
	task="text-classification",
	model=model_name,
	device=device,
	top_k=None # Returns probabilities for all classes
	)
	return MODELS[model_name]


	############################################################
	# Benchmark Examples
	############################################################

	EXAMPLES = {
	"Example 1 (Factual Question)": "India is the most populous country in the world. Do you agree?",
	"Example 2 (Complex Age Calculation)": "Assume that I was born 16 years 4 months 3 days exactly after Sachin Tendulkar was born. How old I am today?",
	"Example 3 (Short Age Query)": "I was born 16 years 4 months 3 days exactly after Sachin Tendulkar was born. My age?",
	"Example 4 (Input-Text Entity Extraction)": "Kalyan KS is from India. Extract named entities from the given sentence.",
	"Example 5 (Task Instruction with Input)": "Extract named entities: Kalyan KS is from India.",
	}


	############################################################
	# UI & Helper Functions
	############################################################

	def update_model_info(selected_model_key):
	url = MODEL_URLS[selected_model_key]
	model_id = MODEL_OPTIONS[selected_model_key]
	perf = MODEL_PERFORMANCE[selected_model_key]

	info = f"""
	- Model: [{model_id}]({url})
	- Architecture: ModernBERT Encoder
	- Parameters: {selected_model_key.split('(')[-1].replace(')', '')}
	- Task: Fact-Check Need Classification (Single-Label)
	- Classes: `FACT_CHECK_NEEDED`, `NO_FACT_CHECK_NEEDED`
	- Performance: {perf}
	"""
	footer = f"""
	### 🔗 Resources & Links
	- Hugging Face Model: [{model_id}]({url})
	- Creator: [Kalyan KS](https://huggingface.co/kalyan-ks)
	- License: Apache 2.0

	---
	Built with ❤️ using Gradio and Hugging Face Transformers.
	"""
	return info, footer


	def update_stats(text):
	if not text:
	return "Characters: 0", "Approx. Tokens: 0"
	char_count = len(text)
	approx_tokens = int(len(text.split()) * 1.3)
	return f"Characters: {char_count}", f"Approx. Tokens: {approx_tokens}"


	def load_example_text(example_key):
	return EXAMPLES.get(example_key, "")


	def clear_inputs():
	empty_df = pd.DataFrame(columns=["Class Label", "Probability", "Confidence Bar"])
	return (
	"", # input_text
	"Characters: 0", # char_caption
	"Approx. Tokens: 0", # token_caption
	"", # status_output
	"", # result_card_html
	empty_df, # prob_table
	"", # metric_prediction
	"", # metric_confidence
	"", # metric_routing
	)


	@spaces.GPU
	def process_classification(input_text, selected_model_key):
	if not input_text or not input_text.strip():
	empty_df = pd.DataFrame(columns=["Class Label", "Probability", "Confidence Bar"])
	return (
	"⚠️ Please enter a prompt to classify.",
	"",
	empty_df,
	"-",
	"-",
	"-",
	)

	model_id = MODEL_OPTIONS[selected_model_key]
	classifier = get_factcheck_pipeline(model_id)

	# Run inference
	raw_results = classifier(input_text.strip())

	if isinstance(raw_results, list) and len(raw_results) > 0 and isinstance(raw_results[0], list):
	scores_list = raw_results[0]
	elif isinstance(raw_results, list):
	scores_list = raw_results
	else:
	scores_list = [raw_results]

	scores_dict = {item["label"]: float(item["score"]) for item in scores_list}

	top_label = max(scores_dict, key=scores_dict.get)
	top_score = scores_dict[top_label]

	# Styling based on prediction
	if top_label == "FACT_CHECK_NEEDED":
	badge_bg = "#FFEBEE"
	badge_border = "#EF5350"
	badge_color = "#C62828"
	icon = "⚠️"
	action_title = "FACT CHECK NEEDED"
	routing_desc = "Route to Retrieval (RAG) / Search / Knowledge Base before generating response."
	routing_short = "Trigger RAG / Web Search"
	else:
	badge_bg = "#E8F5E9"
	badge_border = "#66BB6A"
	badge_color = "#2E7D32"
	icon = "✅"
	action_title = "NO FACT CHECK NEEDED"
	routing_desc = "Safe for direct LLM generation (Creative writing, code, math logic, or input-grounded task)."
	routing_short = "Direct LLM Generation"

	# HTML Result Card
	styled_html = f"""
	<div style="
	border: 2px solid {badge_border};
	background: {badge_bg};
	border-radius: 12px;
	padding: 20px;
	margin-top: 10px;
	box-shadow: 0 2px 8px rgba(0,0,0,0.05);">
	<div style="display: flex; align-items: center; justify-content: space-between; flex-wrap: wrap; gap: 10px;">
	<div style="display: flex; align-items: center; gap: 12px;">
	<span style="font-size: 28px;">{icon}</span>
	<div>
	<h3 style="margin: 0; color: {badge_color}; font-size: 20px; font-weight: 700; letter-spacing: 0.5px;">
	{action_title}
	</h3>
	<p style="margin: 4px 0 0 0; color: #424242; font-size: 14px;">
	{routing_desc}
	</p>
	</div>
	</div>
	<div style="text-align: right; min-width: 120px;">
	<span style="font-size: 26px; font-weight: 800; color: {badge_color};">
	{top_score * 100:.2f}%
	</span>
	<div style="font-size: 11px; text-transform: uppercase; color: #616161; font-weight: 600;">
	Confidence
	</div>
	</div>
	</div>
	</div>
	"""

	# Probability Table
	prob_rows = []
	for label, score in sorted(scores_dict.items(), key=lambda x: x[1], reverse=True):
	bar_len = int(score * 20)
	prob_rows.append({
	"Class Label": label,
	"Probability (%)": f"{score * 100:.2f}%",
	"Confidence Bar": "█" * bar_len + "░" * (20 - bar_len),
	})
	prob_df = pd.DataFrame(prob_rows)

	status_msg = f"✅ Classification completed using {selected_model_key}."

	return (
	status_msg,
	styled_html,
	prob_df,
	top_label,
	f"{top_score * 100:.2f}%",
	routing_short,
	)


	############################################################
	# Gradio Interface
	############################################################

	css = """
	#root, body, html, .gradio-container {
	overflow-y: auto !important;
	height: 100vh !important;
	}
	textarea { font-size: 15px !important; }
	"""

	with gr.Blocks(title="🔎 Fact-Check Classifier Demo", css=css) as demo:

	gr.Markdown("# 🔎 Fact-Check Need Classification using Ammu-1.0")
	gr.Markdown(
	"Determine whether a prompt requires factual verification (RAG / Web Search) or can be processed directly "
	"(Creative Writing, Reasoning, Code, or Input-Grounded Tasks) using the fine-tuned Ammu-1.0 FactCheck encoder models."
	)

	with gr.Sidebar():
	gr.Markdown("## ⚙️ Settings")
	selected_model = gr.Dropdown(
	choices=list(MODEL_OPTIONS.keys()),
	value="Ammu-1.0-Base (150M)",
	label="Select Model"
	)

	gr.Markdown("---")
	gr.Markdown("### 📌 Model Information")
	model_info_md = gr.Markdown()

	with gr.Column():
	gr.Markdown("### 📝 Input Prompt")

	with gr.Row():
	example_selector = gr.Dropdown(
	choices=list(EXAMPLES.keys()),
	value="Example 1 (Factual Question)",
	label="Choose Example"
	)
	load_example_btn = gr.Button("📄 Load Example")

	input_text = gr.Textbox(
	value=EXAMPLES["Example 1 (Factual Question)"],
	lines=5,
	placeholder="Type or paste a user prompt to classify...",
	label="Enter Prompt"
	)

	with gr.Row():
	char_caption = gr.Markdown("Characters: 0")
	token_caption = gr.Markdown("Approx. Tokens: 0")

	with gr.Row():
	classify_btn = gr.Button("🔍 Classify Prompt", variant="primary")
	clear_btn = gr.Button("🗑️ Clear")

	status_output = gr.Markdown()

	with gr.Column():
	gr.Markdown("---")
	gr.Markdown("### 🎯 Classification Result")
	result_card_html = gr.HTML()

	with gr.Row():
	metric_prediction = gr.Textbox(label="Predicted Class", interactive=False)
	metric_confidence = gr.Textbox(label="Confidence", interactive=False)
	metric_routing = gr.Textbox(label="Routing Recommendation", interactive=False)

	gr.Markdown("### 📊 Class Probability Breakdown")
	prob_table = gr.Dataframe(interactive=False)

	gr.Markdown("---")
	footer_md = gr.Markdown()

	# Event Connections
	demo.load(
	fn=update_model_info,
	inputs=[selected_model],
	outputs=[model_info_md, footer_md]
	)

	demo.load(
	fn=update_stats,
	inputs=[input_text],
	outputs=[char_caption, token_caption]
	)

	selected_model.change(
	fn=update_model_info,
	inputs=[selected_model],
	outputs=[model_info_md, footer_md]
	)

	input_text.change(
	fn=update_stats,
	inputs=[input_text],
	outputs=[char_caption, token_caption]
	)

	load_example_btn.click(
	fn=load_example_text,
	inputs=[example_selector],
	outputs=[input_text]
	)

	clear_btn.click(
	fn=clear_inputs,
	outputs=[
	input_text,
	char_caption,
	token_caption,
	status_output,
	result_card_html,
	prob_table,
	metric_prediction,
	metric_confidence,
	metric_routing,
	]
	)

	classify_btn.click(
	fn=process_classification,
	inputs=[input_text, selected_model],
	outputs=[
	status_output,
	result_card_html,
	prob_table,
	metric_prediction,
	metric_confidence,
	metric_routing,
	]
	)

	if __name__ == "__main__":
	demo.launch()