app.py

import math, json
import gradio as gr
import pandas as pd
from typing import Dict, Any
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

LLM_ID = "HuggingFaceTB/SmolLM2-135M-Instruct"
tokenizer = AutoTokenizer.from_pretrained(LLM_ID)
llm = pipeline(
    task="text-generation",
    model=AutoModelForCausalLM.from_pretrained(LLM_ID),
    tokenizer=tokenizer,
    device_map="auto",
)

def projectile_calc(v0_mps: float, theta_deg: float, y0_m: float, g: float, weight_kg: float) -> Dict[str, Any]:
    errors = []
    if not (0 < v0_mps <= 500): errors.append("Initial speed must be in (0, 500] m/s.")
    if not (-10 <= theta_deg <= 90): errors.append("Launch angle must be between -10° and 90°.")
    if not (0 <= y0_m <= 1000): errors.append("Initial height must be in [0, 1000] m.")
    if not (1 <= g <= 50): errors.append("Gravity must be in [1, 50] m/s².")
    if not (0.05 <= weight_kg <= 50): errors.append("Ball weight must be in [0.05, 50] kg.")
    if errors:
        return {"ok": False, "errors": errors}

    theta_rad = math.radians(theta_deg)
    v0x = v0_mps * math.cos(theta_rad)
    v0y = v0_mps * math.sin(theta_rad)

    disc = v0y**2 + 2.0 * g * y0_m
    t_flight = (v0y + math.sqrt(max(disc, 0.0))) / g
    t_flight = max(t_flight, 0.0)

    range_m = v0x * t_flight
    t_apex = max(v0y / g, 0.0)
    y_apex = y0_m + v0y * t_apex - 0.5 * g * t_apex**2
    vy_final = -math.sqrt(max(v0y**2 + 2.0 * g * y0_m, 0.0))
    v_impact = math.sqrt(v0x**2 + vy_final**2)

    return {
        "ok": True,
        "inputs": {
            "v0_mps": v0_mps,
            "theta_deg": theta_deg,
            "y0_m": y0_m,
            "g_mps2": g,
            "weight_kg": weight_kg,
        },
        "derived": {
            "t_apex_s": t_apex,
        },
        "outputs": {
            "time_of_flight_s": t_flight,
            "range_m": range_m,
            "max_height_m": y_apex,
            "impact_speed_mps": v_impact,
        },
        "notes": {
            "assumptions": ["No air resistance.", "Flat landing at y=0.", "Constant gravity g."],
        }
    }

SYSTEM_PROMPT = (
    "You write a single concise sentence using ONLY numbers provided in the JSON. "
    "Do not invent numbers or facts. Use units exactly as given; round to 2–3 sig figs if needed."
)

def llm_one_sentence(structured: Dict[str, Any]) -> str:
    i = structured["inputs"]
    d = structured["derived"]
    o = structured["outputs"]
    payload = {
        "inputs": {
            "v0_mps": round(i["v0_mps"], 3),
            "theta_deg": round(i["theta_deg"], 3),
            "y0_m": round(i["y0_m"], 3),
            "g_mps2": round(i["g_mps2"], 3),
            "weight_kg": round(i["weight_kg"], 3),
        },
        "derived": {
            "t_apex_s": round(d["t_apex_s"], 4),
        },
        "outputs": {
            "time_of_flight_s": round(o["time_of_flight_s"], 4),
            "range_m": round(o["range_m"], 3),
            "max_height_m": round(o["max_height_m"], 3),
            "impact_speed_mps": round(o["impact_speed_mps"], 3),
        }
    }
    instruction = (
        "Produce EXACTLY one sentence of this form using only the JSON values:\n"
        ""If you threw a ball that weighed {weight_kg} kg at v0={v0_mps} m/s, "
        "θ={theta_deg}°, from y0={y0_m} m under g={g_mps2} m/s², "
        "it would stay in the air for {time_of_flight_s} s, reach {max_height_m} m, "
        "travel {range_m} m, and impact at {impact_speed_mps} m/s.""
    )
    user_prompt = f"JSON:\n{json.dumps(payload, indent=2)}\n\nFollow the instruction exactly."
    prompt = f"<|system|>\n{SYSTEM_PROMPT}\n<|user|>\n{instruction}\n\n{user_prompt}\n<|assistant|>"
    out = llm(prompt, max_new_tokens=120, do_sample=False, temperature=0.0, return_full_text=False)
    return out[0]["generated_text"].strip()

def run_once(v0_mps, theta_deg, y0_m, g_mps2, weight_kg):
    rec = projectile_calc(float(v0_mps), float(theta_deg), float(y0_m), float(g_mps2), float(weight_kg))
    if not rec.get("ok", False):
        errs = rec.get("errors", ["Invalid input."])
        df = pd.DataFrame([{"Error": "; ".join(errs)}])
        return df, "Please adjust inputs."
    i, d, o = rec["inputs"], rec["derived"], rec["outputs"]
    df = pd.DataFrame([{
        "v0_mps [m/s]": round(i["v0_mps"], 3),
        "theta_deg [deg]": round(i["theta_deg"], 3),
        "y0_m [m]": round(i["y0_m"], 3),
        "g_mps2 [m/s²]": round(i["g_mps2"], 3),
        "weight_kg [kg]": round(i["weight_kg"], 3),
        "t_apex_s [s]": round(d["t_apex_s"], 4),
        "max_height_m [m]": round(o["max_height_m"], 3),
        "time_of_flight_s [s]": round(o["time_of_flight_s"], 4),
        "range_m [m]": round(o["range_m"], 3),
        "impact_speed_mps [m/s]": round(o["impact_speed_mps"], 3),
    }])
    sentence = llm_one_sentence(rec)
    return df, sentence

with gr.Blocks(title="Projectile Motion — Deterministic + One-Sentence LLM") as demo:
    gr.Markdown("# Projectile Motion — Deterministic Calculator + One-Sentence Explanation")
    gr.Markdown("Deterministic projectile motion (no air resistance). See all inputs and derived values, plus a single, grounded sentence.")

    with gr.Row():
        v0 = gr.Slider(1, 200, value=30.0, step=0.5, label="Initial speed v₀ [m/s]")
        theta = gr.Slider(-10, 90, value=45.0, step=0.5, label="Launch angle θ [deg]")
        y0 = gr.Slider(0, 20, value=1.5, step=0.1, label="Initial height y₀ [m]")
        g = gr.Slider(5, 20, value=9.81, step=0.01, label="Gravity g [m/s²]")
        w = gr.Slider(0.05, 10.0, value=0.43, step=0.01, label="Ball weight [kg]")

    run_btn = gr.Button("Compute")

    results_df = gr.Dataframe(
        headers=[
            "v0_mps [m/s]", "theta_deg [deg]", "y0_m [m]", "g_mps2 [m/s²]", "weight_kg [kg]",
            "t_apex_s [s]", "max_height_m [m]", "time_of_flight_s [s]", "range_m [m]", "impact_speed_mps [m/s]"
        ],
        label="All values (inputs + derived)",
        interactive=False
    )
    explain_md = gr.Markdown(label="One-sentence explanation")

    run_btn.click(run_once, inputs=[v0, theta, y0, g, w], outputs=[results_df, explain_md])

    gr.Examples(
        examples=[
            [30.0, 45.0, 1.5, 9.81, 0.43],
            [20.0, 30.0, 0.0, 9.81, 0.145],
            [40.0, 60.0, 0.0, 9.81, 0.45],
        ],
        inputs=[v0, theta, y0, g, w],
        label="Representative cases",
        examples_per_page=3,
        cache_examples=False,
    )

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