app.py

import gradio as gr
from transformers import AutoModel
import torch
import numpy as np

tokenizer = torch.hub.load('huggingface/pytorch-transformers', 'tokenizer', 'bert-base-cased-finetuned-mrpc')
model = AutoModel.from_pretrained("jacklindsai/is_it_elon_musk")

def preprocess_text(text):
    return tokenizer.encode_plus(text, truncation=True, padding='max_length', max_length=48, return_attention_mask=True)

device = torch.device('cpu')
def pred_is_elon_musk(text):
    encoded_text = preprocess_text(text)
    ids = encoded_text['input_ids']
    masks = encoded_text['attention_mask']
    ids = torch.Tensor(ids).to(device, dtype=torch.int32)
    masks = torch.Tensor(masks).to(device, dtype=torch.int32)
    results = model(input_ids=ids, token_type_ids=None,
                    attention_mask=masks)
    logis = results['logits'].detach().numpy()
    prediction = np.argmax(logis, axis=1).flatten()
    if prediction == 1:
        return "It's from Elon Musk."
    else:
        return "It's NOT from Elon Musk."

iface = gr.Interface(pred_is_elon_musk, inputs="text", outputs="text", title='Elon Musk Classifier APP', theme = "dark-peach", examples=["Now I'm going to buy McDonald's and fix all the ice cream machines...","Next I’m buying Coca-Cola to put the cocaine back in"], description="This app predicts whether the tweet is from Elon Musk.")

iface.launch(inline=False)