app.py

import math # Used for mathematical functions like radians and pi
import gradio as gr # Framework for building the web interface
import pandas as pd # Used for creating the structured output table (DataFrame)

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline # Core components for running the HuggingFace LLM


# Specifies the smaller, instructional model for low-latency recommendations
MODEL_ID = "HuggingFaceTB/SmolLM2-135M-Instruct"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
pipe = pipeline(
    task="text-generation",
    model=AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
    ),
    tokenizer=tokenizer,
    device=-1 # Runs on CPU by default; set device=0 for GPU
)

# 1. Calculation Function

def gear_calc(N_teeth: int, Pd: float, phi_deg: float) -> dict:
    """
    Calculates standard dimensions for an external spur gear based on common engineering formulas.

    Inputs:
    - N_teeth (int): Number of teeth
    - Pd (float): Diametral Pitch [teeth/inch]
    - phi_deg (float): Pressure Angle in degrees

    Returns:
    - dict: A dictionary containing all standard gear dimensions in both mm and inches.
    """
    if N_teeth <= 0 or Pd <= 0 or phi_deg <= 0:
        # Input validation: all gear parameters must be positive
        raise ValueError("All inputs must be positive.")

    phi_rad = math.radians(phi_deg)

    # Pitch Diameter (D) - The imaginary circle upon which the tooth spacing is measured.
    pitch_diameter_in = N_teeth / Pd
    pitch_diameter_mm = pitch_diameter_in * 25.4

    # Addendum (a) and Dedendum (b) - Heights above and depths below the pitch circle.
    addendum_in = 1 / Pd
    dedendum_in = 1.25 / Pd
    addendum_mm = addendum_in * 25.4
    dedendum_mm = dedendum_in * 25.4

    # Outside (OD) and Root Diameters - The largest and smallest diameters of the gear.
    od_in = pitch_diameter_in + 2 * addendum_in
    root_diameter_in = pitch_diameter_in - 2 * dedendum_in
    od_mm = od_in * 25.4
    root_diameter_mm = root_diameter_in * 25.4

    # Working Depth and Whole Depth - Total contact depth and total tooth height.
    working_depth_in = 2 * addendum_in
    whole_depth_in = addendum_in + dedendum_in
    working_depth_mm = working_depth_in * 25.4
    whole_depth_mm = whole_depth_in * 25.4

    # Circular Pitch (p) and Tooth Thickness (t) - Spacing and width of the tooth along the pitch circle.
    circular_pitch_in = math.pi / Pd
    circular_pitch_mm = circular_pitch_in * 25.4
    tooth_thickness_in = circular_pitch_in / 2
    tooth_thickness_mm = tooth_thickness_in * 25.4

    # Base Diameter (Db) - The circle from which the involute curve is generated.
    base_diameter_in = pitch_diameter_in * math.cos(phi_rad)
    base_diameter_mm = base_diameter_in * 25.4

    # Return a comprehensive dictionary of all calculated dimensions
    return dict(
        pitch_diameter_mm=pitch_diameter_mm,
        pitch_diameter_in=pitch_diameter_in,
        od_mm=od_mm,
        od_in=od_in,
        root_diameter_mm=root_diameter_mm,
        root_diameter_in=root_diameter_in,
        addendum_mm=addendum_mm,
        addendum_in=addendum_in,
        dedendum_mm=dedendum_mm,
        dedendum_in=dedendum_in,
        working_depth_mm=working_depth_mm,
        working_depth_in=working_depth_in,
        whole_depth_mm=whole_depth_mm,
        whole_depth_in=whole_depth_in,
        circular_pitch_mm=circular_pitch_mm,
        circular_pitch_in=circular_pitch_in,
        tooth_thickness_mm=tooth_thickness_mm,
        tooth_thickness_in=tooth_thickness_in,
        base_diameter_mm=base_diameter_mm,
        base_diameter_in=base_diameter_in,
    )


# 2. LLM Helper Functions

def _format_chat(system_prompt: str, user_prompt: str) -> str:
    """
    Helper function to structure the system and user prompts into the LLM's required
    chat template format for optimal instruction adherence.
    """
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]
    template = getattr(tokenizer, "chat_template", None)
    return tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

def _llm_generate(prompt: str, max_tokens: int) -> str:
    """
    Executes the text generation pipeline using deterministic settings (do_sample=False,
    temperature=0.0) to ensure the LLM provides the most compliant, non-creative output.
    """
    out = pipe(
        prompt,
        max_new_tokens=max_tokens,
        do_sample=False,   # Disable random sampling for deterministic results
        temperature=0.0,   # Set temperature to 0.0 for maximum compliance
        return_full_text=False,
    )
    # Strip leading/trailing whitespace from the generated text
    return out[0]["generated_text"].strip()


# 3. LLM Explanation Function

def llm_explain(results: dict, inputs: list) -> str:
    """
    Generates a clear, customized explanation for a non-technical user
    based on the calculated gear dimensions.
    """

    pitch_diameter_in = results['pitch_diameter_in']
    tooth_thickness_in = results['tooth_thickness_in']
    circular_pitch_in = results['circular_pitch_in']

    # Estimate scale
    if pitch_diameter_in < 2:
        scale = "small"
        examples = "clocks, cameras, or precision instruments"
    elif pitch_diameter_in < 10:
        scale = "medium"
        examples = "bicycles, sewing machines, or automotive gearboxes"
    else:
        scale = "large"
        examples = "wind turbines, cranes, or heavy mining machinery"

    # Interpret tooth strength
    if tooth_thickness_in < 0.05:
        strength = "delicate and precise, best for very light loads"
    elif tooth_thickness_in < 0.3:
        strength = "balanced between precision and strength"
    else:
        strength = "very strong, made for heavy-duty work"

    # Interpret circular pitch (spacing)
    if circular_pitch_in < 0.2:
        spacing = "fine and closely spaced for smooth, quiet motion"
    elif circular_pitch_in < 1:
        spacing = "moderate spacing for reliable general use"
    else:
        spacing = "wide spacing, rugged enough to handle shocks and dirt"

    # System prompt
    system_prompt = (
        "You are a mechanical engineer explaining gears to a non-technical user. "
        "Write in 2–4 friendly sentences. "
        "Explain what the size suggests, what the tooth shape/spacing means in practice, "
        "and give unique example applications that fit the gear’s scale. "
        "Do not simply restate the numbers given. "
        "Be specific and avoid repeating the same examples across different gear sizes."
    )

    # User prompt (gear context)
    user_prompt = (
        f"This is a {scale} gear. "
        f"Its teeth are {strength}, and the spacing is {spacing}. "
        f"Suggest everyday uses where a gear like this would make sense, such as {examples}. "
        "Explain it in a way that feels useful and practical for someone who is not an engineer."
    )

    formatted = _format_chat(system_prompt, user_prompt)

    explanation = _llm_generate(formatted, max_tokens=200)

    # Fallback if the LLM fails
    if not explanation:
        explanation = (
            f"This is a {scale} gear with {strength} and {spacing}. "
            f"You’d expect to find gears like this in {examples}."
        )

    return explanation


# 4. Main Entry Point for GUI

def run_once(N_teeth_str, Pd_str, phi_deg_str):
    """
    Main function executed by the Gradio interface. Handles input parsing,
    calculation, result formatting (DataFrame), and LLM generation.
    """
    try:
        # Input conversion from string to required types
        N_teeth = int(N_teeth_str)
        Pd = float(Pd_str)
        phi_deg = float(phi_deg_str)
        inputs = [N_teeth, Pd, phi_deg]

        # Perform the core mechanical calculation
        results = gear_calc(N_teeth, Pd, phi_deg)

        # Prepare the data structure for the Pandas DataFrame output table
        rows = [
            # Each dictionary represents a row in the output table
            {"Quantity": "Pitch Diameter", "Value (mm)": results["pitch_diameter_mm"], "Value (in)": results["pitch_diameter_in"]},
            {"Quantity": "Outside Diameter (OD)", "Value (mm)": results["od_mm"], "Value (in)": results["od_in"]},
            {"Quantity": "Root Diameter", "Value (mm)": results["root_diameter_mm"], "Value (in)": results["root_diameter_in"]},
            {"Quantity": "Base Diameter", "Value (mm)": results["base_diameter_mm"], "Value (in)": results["base_diameter_in"]},
            {"Quantity": "Addendum", "Value (mm)": results["addendum_mm"], "Value (in)": results["addendum_in"]},
            {"Quantity": "Dedendum", "Value (mm)": results["dedendum_mm"], "Value (in)": results["dedendum_in"]},
            {"Quantity": "Working Depth", "Value (mm)": results["working_depth_mm"], "Value (in)": results["working_depth_in"]},
            {"Quantity": "Whole Depth", "Value (mm)": results["whole_depth_mm"], "Value (in)": results["whole_depth_in"]},
            {"Quantity": "Circular Pitch", "Value (mm)": results["circular_pitch_mm"], "Value (in)": results["circular_pitch_in"]},
            {"Quantity": "Tooth Thickness", "Value (mm)": results["tooth_thickness_mm"], "Value (in)": results["tooth_thickness_in"]},
        ]

        # Create and format the output DataFrame, rounding values for display
        df = pd.DataFrame(rows)
        df["Value (mm)"] = df["Value (mm)"].round(4)
        df["Value (in)"] = df["Value (in)"].round(4)

        # Generate the LLM narrative
        narrative = llm_explain(results, inputs)
        
        return df, narrative

    except ValueError as e:
        # Handle mathematical and input validation errors
        error_df = pd.DataFrame([{"Error": str(e)}])
        return error_df, "Error: Please ensure all inputs are valid positive numbers. Teeth must be an integer."
    except Exception as e:
        # Handle unexpected system/LLM errors
        error_df = pd.DataFrame([{"System Error": "Calculation or LLM failed"}])
        print(f"Internal Error: {e}")
        return error_df, "An unexpected system error occurred during computation or LLM generation."


# Gradio UI Definition

with gr.Blocks() as demo:

    # App Title and Description
    gr.Markdown(
        "# External Gear Dimension Calculator ⚙️"
    )
    gr.Markdown(
        "Enter number of teeth, diametral pitch, and pressure angle to get standard external gear dimensions, and an LLM-generated use case."
    )

    # Input parameters organized in a row layout
    with gr.Row():
        N_teeth_in = gr.Number(value=0, label="Number of Teeth (#)", precision=0)
        Pd_in = gr.Number(value=0, label="Diametral Pitch [teeth/inch]")
        phi_deg_in = gr.Number(value=0, label="Pressure Angle [°]")

    # Action button to trigger the process
    run_btn = gr.Button("Calculate and Explain Gear Use")

    # Output components: DataFrame for numbers, Markdown for LLM text
    results_df = gr.Dataframe(label="Gear Dimensions (mm and inches)", interactive=False)
    explain_md = gr.Markdown(label="LLM-Generated Use Case")

    # Event listener: Connect the button click to the main computation function
    run_btn.click(fn=run_once, inputs=[N_teeth_in, Pd_in, phi_deg_in], outputs=[results_df, explain_md])

    # Pre-defined examples for easy testing
    gr.Examples(
        examples=[
            [12, 4.0, 20.0],   
            [40, 0.75, 45.0],  
            [200, 0.5, 25.0],  
        ],
        inputs=[N_teeth_in, Pd_in, phi_deg_in],
        label="Representative Gear Cases",
        examples_per_page=3,
        cache_examples=False,
    )

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