Karenuppp/TextClassification

1import string2import gradio as gr3import requests4import torch5from transformers import (6    AutoConfig,7    AutoModelForSequenceClassification,8    AutoTokenizer,9)10 11model_dir = "my-bert-model"12 13config = AutoConfig.from_pretrained(model_dir, num_labels=3, finetuning_task="text-classification")14tokenizer = AutoTokenizer.from_pretrained(model_dir)15model = AutoModelForSequenceClassification.from_pretrained(model_dir, config=config)16 17def inference(input_text):18    inputs = tokenizer.batch_encode_plus(19                [input_text],20                max_length=512,21                pad_to_max_length=True,22                truncation=True,23                padding="max_length",24                return_tensors="pt",25            )26    27    with torch.no_grad():28        logits = model(**inputs).logits29    30    predicted_class_id = logits.argmax().item()31    output = model.config.id2label[predicted_class_id]32    return output33 34with gr.Blocks(css="""35    .message.svelte-w6rprc,svelte-w6rprc,svelte-w6rprc {font-size: 20px; margin-top: 20px}36    #component-21 > div.wrap.svelte-w6rprc {height: 600px;}37    """) as demo:38    with gr.Row():39        with gr.Column():40            input_text = gr.Textbox(placeholder="Insert your prompt here:", scale=2, container=False)41            answer = gr.Textbox(lines=0,label="Answer")42            generate_bt = gr.Button("Generate",scale=1)43    inputs =[input_text]44    outputs = [answer]45    generate_bt.click(46        fn=inference, inputs=inputs, outputs=outputs, show_progress=True47    )48    examples =[49        ["My last two weather pics from the storm on August 2nd. people packed up real fast after the temp dropped and winds picked up.", 1],50        ["Lying Clinton sinking! Donald Trump singing: Let's Make America Great Again!", 0],51        ],52demo.queue()53demo.launch()54