rifatramadhani/topic-classification

1import gradio as gr2import spaces3import torch4from transformers import pipeline5import datetime6import json7import logging8 9model_path = "cardiffnlp/twitter-roberta-base-dec2021-tweet-topic-multi-all"10# Load model for first time cache 11topic_classification_task = pipeline("text-classification", model=model_path, tokenizer=model_path)12 13@spaces.GPU14def classify(query):15    torch_device = 0 if torch.cuda.is_available() else -116    tokenizer_kwargs = {'truncation':True,'max_length':512}17 18    topic_classification_task = pipeline("text-classification", model=model_path, tokenizer=model_path, device=torch_device)19 20    request_type = type(query)21    try:22        data = json.loads(query)23        if type(data) != list:24            data = [query]25        else:26            request_type = type(data)27    except Exception as e:28        print(e)29        data = [query]30        pass31 32    start_time = datetime.datetime.now()33    34    result = topic_classification_task(data, batch_size=128, top_k=3, **tokenizer_kwargs)35 36    end_time = datetime.datetime.now()37    elapsed_time = end_time - start_time38 39    logging.debug("elapsed predict time: %s", str(elapsed_time))40    print("elapsed predict time:", str(elapsed_time))41    42    output = {}43    output["time"] = str(elapsed_time)44    output["device"] = torch_device45    output["result"] = result46 47    return json.dumps(output)48 49demo = gr.Interface(fn=classify, inputs=["text"], outputs="text")50demo.launch()