Add Roberta model

app.py

@@ -24,6 +24,10 @@ models = {
     "TinyBERT": {
         "tokenizer": AutoTokenizer.from_pretrained("elo4/TinyBERT-sentiment-model"),
         "model": AutoModelForSequenceClassification.from_pretrained("elo4/TinyBERT-sentiment-model"),
+    },
+    "RoBERTa": {
+        "tokenizer": AutoTokenizer.from_pretrained("ordek899/roberta_1to5rating_pred_for_restaur_trained_on_hotels"),
+        "model": AutoModelForSequenceClassification.from_pretrained("ordek899/roberta_1to5rating_pred_for_restaur_trained_on_hotels"),
     }
 }
 
@@ -82,6 +86,16 @@ def predict_with_tinybert(text):
         predictions = logits.argmax(axis=-1).cpu().numpy()
     return int(predictions[0] + 1)
 
+def predict_with_roberta_ordek899(text):
+    tokenizer = models["RoBERTa"]["tokenizer"]
+    model = models["RoBERTa"]["model"]
+    encodings = tokenizer([text], padding=True, truncation=True, max_length=512, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**encodings)
+        logits = outputs.logits
+        predictions = logits.argmax(axis=-1).cpu().numpy()
+    return int(predictions[0] + 1)
+
 # Unified function for sentiment analysis and statistics
 def analyze_sentiment_and_statistics(text):
     results = {
@@ -89,6 +103,7 @@ def analyze_sentiment_and_statistics(text):
         "Logistic Regression": predict_with_logistic_regression(text),
         "BERT Multilingual (NLP Town)": predict_with_bert_multilingual(text),
         "TinyBERT": predict_with_tinybert(text),
+        "RoBERTa": predict_with_roberta_ordek899(text),
     }
     
     # Calculate statistics
@@ -156,6 +171,7 @@ with gr.Blocks(css=".gradio-container { max-width: 900px; margin: auto; padding:
             log_reg_output = gr.Textbox(label="Predicted Sentiment (Logistic Regression)", interactive=False)
             bert_output = gr.Textbox(label="Predicted Sentiment (BERT Multilingual)", interactive=False)
             tinybert_output = gr.Textbox(label="Predicted Sentiment (TinyBERT)", interactive=False)
+            roberta_ordek_output = gr.Textbox(label="Predicted Sentiment (RoBERTa)", interactive=False)
         
         with gr.Column():
             statistics_output = gr.Textbox(label="Statistics (Lowest, Highest, Average)", interactive=False)
@@ -169,6 +185,7 @@ with gr.Blocks(css=".gradio-container { max-width: 900px; margin: auto; padding:
                 f"{results['Logistic Regression']}",
                 f"{results['BERT Multilingual (NLP Town)']}",
                 f"{results['TinyBERT']}",
+                f"{results['RoBERTa']}",
                 f"Statistics:\n{statistics['Message']}\nAverage Score: {statistics['Average Score']}"
             )
         else:  # Min and Max scores are present
@@ -177,13 +194,21 @@ with gr.Blocks(css=".gradio-container { max-width: 900px; margin: auto; padding:
                 f"{results['Logistic Regression']}",
                 f"{results['BERT Multilingual (NLP Town)']}",
                 f"{results['TinyBERT']}",
+                f"{results['RoBERTa']}",
                 f"Statistics:\n{statistics['Lowest Score']}\n{statistics['Highest Score']}\nAverage Score: {statistics['Average Score']}"
             )
     
     analyze_button.click(
         process_input_and_analyze,
         inputs=[text_input],
-        outputs=[distilbert_output, log_reg_output, bert_output, tinybert_output, statistics_output]
+        outputs=[
+            distilbert_output, 
+            log_reg_output, 
+            bert_output, 
+            tinybert_output, 
+            roberta_ordek_output, 
+            statistics_output
+        ]
     )
 
 # Launch the app