import datasets import numpy as np import torch import transformers from config import epochs, batch_size, learning_rate, id2label from model import tokenizer, multitask_model from mtm import MultitaskTrainer, NLPDataCollator, DataLoaderWithTaskname import pandas as pd from datasets import Dataset, DatasetDict from data_predict import convert_to_stsb_features,convert_to_features import gradio as gr from huggingface_hub import hf_hub_download,snapshot_download device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_link = hf_hub_download(repo_id="FFZG-cleopatra/Croatian-News-Classifier",filename = "pytorch_model.bin") multitask_model.load_state_dict(torch.load(model_link, map_location=device)) multitask_model.to(device) def predict_sentiment(sentence = "Volim ti"): # gather everyone if you want to have a single DatasetDict document = DatasetDict({ # "train": Dataset.from_pandas(df_document_sl_hr_train), # "valid": Dataset.from_pandas(df_document_sl_hr_valid), "test": Dataset.from_dict({"content":[sentence]}) }) dataset_dict = { "document": document, } for task_name, dataset in dataset_dict.items(): print(task_name) print(dataset_dict[task_name]["test"][0]) print() convert_func_dict = { "document": convert_to_stsb_features, # "paragraph": convert_to_stsb_features, # "sentence": convert_to_stsb_features, } features_dict = convert_to_features(dataset_dict, convert_func_dict) predictions = [] for _, batch in enumerate(features_dict["document"]['test']): for key, value in batch.items(): batch[key] = batch[key].to(device) task_model = multitask_model.get_model("document") classifier_output = task_model.forward( torch.unsqueeze(batch["input_ids"], 0), torch.unsqueeze(batch["attention_mask"], 0),) print(tokenizer.decode(batch["input_ids"],skip_special_tokens=True)) prediction =torch.max(classifier_output.logits, axis=1) predictions.append(prediction.indices.item()) print("p:", predictions[0] , id2label[predictions[0]] ) return id2label[predictions[0]] interface = gr.Interface( fn=predict_sentiment, inputs='text', outputs=['label'], title='Sentiment Analysis', description='Get the positive/neutral/negative sentiment for the given input.' ) interface.launch(inline = False)