Update app.py

app.py

@@ -1,11 +1,30 @@
 import gradio as gr
-from transformers import pipeline
 from gradio.components import Textbox
 
-# Load the sentiment analysis pipeline with DistilBERT
-distilbert_pipeline = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
-label_map = {"POSITIVE":"OTHER", "NEGATIVE":"SENSITIVE"}
+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):
     """
@@ -13,10 +32,8 @@ def predict_sentiment(text):
     :param text: str, input text to analyze.
     :return: str, predicted sentiment and confidence score.
     """
-    result = distilbert_pipeline(text)[0]
-    label = label_map[result['label']]
-    score = result['score']
-    return f"TAG: {label}, Confidence: {score:.2f}"
+    result = is_pharma(text)
+    return f"TAG: {result}" #, Confidence: {score:.2f}
 
 input1 = Textbox(lines=2, placeholder="Type your text here...")
 
@@ -24,8 +41,8 @@ input1 = Textbox(lines=2, placeholder="Type your text here...")
 iface = gr.Interface(fn=predict_sentiment,
                      inputs=input1,
                      outputs="text",
-                     title="Talk2Loop Sensitive statement tags",
-                     description="This model predicts the sensitivity of the input text. Enter a sentence to see if it's sensitive or not.")
+                     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()