Spaces:
Running
Running
File size: 3,747 Bytes
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 142 143 144 145 |
from __future__ import annotations
from pathlib import Path
from typing import Literal
import click
__all__ = ["cli_wrapper"]
ERROR_STR = click.style("ERROR", fg="red")
DONE_STR = click.style("DONE", fg="green")
POSITIVE_STR = click.style("POSITIVE", fg="green")
NEUTRAL_STR = click.style("NEUTRAL", fg="yellow")
NEGATIVE_STR = click.style("NEGATIVE", fg="red")
@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:
click.echo(f"{ERROR_STR}: No text provided")
return
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 = NEGATIVE_STR
elif prediction == 1:
sentiment = POSITIVE_STR
else:
sentiment = NEUTRAL_STR
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()
|