| from typing import Optional |
|
|
| import os |
| import pandas as pd |
| import numpy as np |
| import tabpfn_client |
| from tabpfn_client import TabPFNClassifier |
| from tabpfn_client.constants import ModelVersion |
|
|
|
|
| def init_tabpfn(token: Optional[str] = None) -> None: |
| """Configure tabpfn-client. Pass token explicitly or set TABPFN_API_TOKEN in the environment.""" |
| resolved = (token or "").strip() or os.environ.get("TABPFN_API_TOKEN", "").strip() |
| if not resolved: |
| raise RuntimeError("TabPFN API token is required.") |
| tabpfn_client.set_access_token(resolved) |
|
|
| FEATURES = [ |
| "fixed_acidity", "volatile_acidity", "citric_acid", "residual_sugar", |
| "chlorides", "free_sulfur_dioxide", "total_sulfur_dioxide", "density", |
| "pH", "sulphates", "alcohol", "wine_type", |
| ] |
|
|
| FEATURE_DISPLAY_NAMES = { |
| "fixed_acidity": "Fixed Acidity", |
| "volatile_acidity": "Volatile Acidity", |
| "citric_acid": "Citric Acid", |
| "residual_sugar": "Residual Sugar (g/L)", |
| "chlorides": "Chlorides", |
| "free_sulfur_dioxide": "Free SO₂ (mg/L)", |
| "total_sulfur_dioxide": "Total SO₂ (mg/L)", |
| "density": "Density (g/cm³)", |
| "pH": "pH", |
| "sulphates": "Sulphates", |
| "alcohol": "Alcohol (%)", |
| "wine_type": "Wine Type (0=Red, 1=White)", |
| } |
|
|
| QUALITY_LABELS = {0: "Low (≤5)", 1: "Medium (6)", 2: "High (≥7)"} |
| |
| QUALITY_INDEX_BY_BIN = {0: 4.5, 1: 6.0, 2: 7.5} |
|
|
|
|
| def quality_index_from_probs(probs: np.ndarray) -> float: |
| return float(sum(probs[i] * QUALITY_INDEX_BY_BIN[i] for i in range(len(probs)))) |
|
|
| RED_URL = "https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv" |
| WHITE_URL = "https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-white.csv" |
|
|
|
|
| def bin_quality(q: int) -> int: |
| if q <= 5: |
| return 0 |
| if q == 6: |
| return 1 |
| return 2 |
|
|
|
|
| def load_data() -> pd.DataFrame: |
| red = pd.read_csv(RED_URL, sep=";") |
| red["wine_type"] = 0 |
| white = pd.read_csv(WHITE_URL, sep=";") |
| white["wine_type"] = 1 |
| df = pd.concat([red, white], ignore_index=True) |
| df.columns = [c.strip().replace(" ", "_") for c in df.columns] |
| df["quality_bin"] = df["quality"].apply(bin_quality) |
| return df |
|
|
|
|
| def prepare_split(df: pd.DataFrame, test_size: int = 1000, random_state: int = 42): |
| df = df.sample(frac=1, random_state=random_state).reset_index(drop=True) |
| test_df = df.iloc[:test_size].reset_index(drop=True) |
| train_df = df.iloc[test_size:].reset_index(drop=True) |
| return train_df, test_df |
|
|
|
|
| def fit_model(train_df: pd.DataFrame) -> TabPFNClassifier: |
| X = train_df[FEATURES].values.astype(float) |
| y = train_df["quality_bin"].values |
| clf = TabPFNClassifier.create_default_for_version(ModelVersion.V3) |
| clf.fit(X, y) |
| return clf |
|
|
|
|
| def predict(clf: TabPFNClassifier, features: dict) -> tuple[str, dict, float, int]: |
| """Returns (predicted label, {label: probability}, estimated quality index, class index).""" |
| row = np.array([[features[f] for f in FEATURES]], dtype=float) |
| probs = clf.predict_proba(row)[0] |
| pred_idx = int(np.argmax(probs)) |
| prob_map = {QUALITY_LABELS[i]: float(probs[i]) for i in range(len(probs))} |
| quality_index = round(quality_index_from_probs(probs), 1) |
| return QUALITY_LABELS[pred_idx], prob_map, quality_index, pred_idx |
|
|
