Spaces:
Running
Running
File size: 3,623 Bytes
85ac990 0993d5e 85ac990 0993d5e 85ac990 0993d5e 85ac990 0993d5e 85ac990 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
from __future__ import annotations
from pathlib import Path
from typing import Literal
import click
__all__ = ["cli_wrapper"]
DONE_STR = click.style("DONE", fg="green")
@click.group()
def cli() -> None: ...
@cli.command()
@click.option(
"--model",
"model_path",
required=True,
help="Path to the trained model",
type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True, resolve_path=True, path_type=Path),
)
@click.option(
"--share/--no-share",
default=False,
help="Whether to create a shareable link",
)
def gui(model_path: Path, share: bool) -> None:
"""Launch the Gradio GUI"""
from app.gui import launch_gui
launch_gui(model_path, share)
@cli.command()
@click.option(
"--model",
"model_path",
required=True,
help="Path to the trained model",
type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True, resolve_path=True, path_type=Path),
)
@click.argument("text", nargs=-1)
def predict(model_path: Path, text: list[str]) -> None:
"""Perform sentiment analysis on the provided text.
Note: Piped input takes precedence over the text argument
"""
import sys
import joblib
text = " ".join(text).strip()
if not sys.stdin.isatty():
piped_text = sys.stdin.read().strip()
text = piped_text or text
if not text:
msg = "No text provided"
raise click.UsageError(msg)
click.echo("Loading model... ", nl=False)
model = joblib.load(model_path)
click.echo(DONE_STR)
click.echo("Performing sentiment analysis... ", nl=False)
prediction = model.predict([text])[0]
if prediction == 0:
sentiment = click.style("NEGATIVE", fg="red")
elif prediction == 1:
sentiment = click.style("POSITIVE", fg="green")
else:
sentiment = click.style("NEUTRAL", fg="yellow")
click.echo(sentiment)
@cli.command()
@click.option(
"--dataset",
required=True,
help="Dataset to train the model on",
type=click.Choice(["sentiment140", "amazonreviews", "imdb50k"]),
)
@click.option(
"--max-features",
default=20000,
help="Maximum number of features",
show_default=True,
type=click.IntRange(1, None),
)
@click.option(
"--seed",
default=42,
help="Random seed (-1 for random seed)",
show_default=True,
type=click.IntRange(-1, None),
)
def train(
dataset: Literal["sentiment140", "amazonreviews", "imdb50k"],
max_features: int,
seed: int,
) -> None:
"""Train the model on the provided dataset"""
import joblib
from app.constants import MODELS_DIR
from app.model import create_model, load_data, train_model
model_path = MODELS_DIR / f"{dataset}_tfidf_ft-{max_features}.pkl"
if model_path.exists():
click.confirm(f"Model file '{model_path}' already exists. Overwrite?", abort=True)
click.echo("Preprocessing dataset... ", nl=False)
text_data, label_data = load_data(dataset)
click.echo(DONE_STR)
click.echo("Creating model... ", nl=False)
model = create_model(max_features, seed=None if seed == -1 else seed)
click.echo(DONE_STR)
click.echo("Training model... ", nl=False)
accuracy = train_model(model, text_data, label_data)
joblib.dump(model, model_path)
click.echo(DONE_STR)
click.echo("Model accuracy: ")
click.secho(f"{accuracy:.2%}", fg="blue")
# TODO: Add hyperparameter options
# TODO: Random/grid search for finding best classifier and hyperparameters
def cli_wrapper() -> None:
cli(max_content_width=120)
if __name__ == "__main__":
cli_wrapper()
|