Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from gradio.components import Textbox | |
| import torch | |
| from transformers import pipeline | |
| from transformers import AutoTokenizer | |
| from transformers import AutoModelForSequenceClassification | |
| # Load the DistilBERT tokenizer | |
| tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased') | |
| #Load the model | |
| model = AutoModelForSequenceClassification.from_pretrained("skylord/pharma_classification") | |
| def is_pharma(sentence, tokenize=tokenizer, model=model): | |
| # tokenize the input | |
| inputs = tokenizer(sentence, return_tensors='pt') | |
| # ensure model and inputs are on the same device (GPU) | |
| inputs = {name: tensor.cuda() for name, tensor in inputs.items()} | |
| model = model.cuda() | |
| # get prediction - 2 classes "probabilities" (not really true because they still need to be normalized) | |
| with torch.no_grad(): | |
| predictions = model(**inputs)[0].cpu().numpy() | |
| # get the top prediction class and convert it to its associated label | |
| top_prediction = predictions.argmax().item() | |
| return ds['train'].features['labels'].int2str(top_prediction) | |
| def predict_sentiment(text): | |
| """ | |
| Predicts the sentiment of the input text using DistilBERT. | |
| :param text: str, input text to analyze. | |
| :return: str, predicted sentiment and confidence score. | |
| """ | |
| result = is_pharma(text) | |
| return f"TAG: {result}" #, Confidence: {score:.2f} | |
| input1 = Textbox(lines=2, placeholder="Type your text here...") | |
| # Create a Gradio interface | |
| iface = gr.Interface(fn=predict_sentiment, | |
| inputs=input1, | |
| outputs="text", | |
| title="Identify if the news item is relevant to the pharma industry", | |
| description="This model predicts the tag of the input text. Enter a sentence to see if it's pharma or not. Response is a Yes or a No") | |
| # Launch the interface | |
| iface.launch() | |