bertas defined

app.py

@@ -4,7 +4,7 @@ from gradio_client import utils as client_utils
 import spaces
 import pandas as pd
 import torch
-from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer, AutoModel
+from transformers import AutoModelForSequenceClassification, AutoModelForSeq2SeqLM, AutoTokenizer, AutoModel
 import plotly.graph_objects as go
 import logging
 import io
@@ -22,6 +22,7 @@ import logging
 from typing import List, Set, Tuple
 import asyncio
 
+
 os.environ["ZEROGPU_ENABLED"] = "1"
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
@@ -322,9 +323,7 @@ class EventDetector:
                 
             self.rutranslator = ru_translate_fn
             
-            # Continue with other models but use safetensors
-            # For sentiment models
-            from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
             
             # Initialize FinBERT
             finbert_tokenizer = AutoTokenizer.from_pretrained("ProsusAI/finbert")
@@ -345,7 +344,44 @@ class EventDetector:
                 
             self.finbert = lambda text: [finbert_sentiment(text)]
             
-            # Similar custom implementations for other sentiment models...
+            # Initialize RoBERTa
+            roberta_tokenizer = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
+            roberta_model = AutoModelForSequenceClassification.from_pretrained(
+                "cardiffnlp/twitter-roberta-base-sentiment", 
+                use_safetensors=True,
+                device_map=device
+            )
+            
+            # Define custom pipeline instead of using transformers pipeline
+            def roberta_sentiment(text):
+                inputs = roberta_tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
+                outputs = roberta_model(**inputs)
+                probs = outputs.logits.softmax(dim=1)[0]
+                pred_idx = probs.argmax().item()
+                labels = ["negative", "neutral", "positive"]
+                return {"label": labels[pred_idx], "score": probs[pred_idx].item()}
+                
+            self.roberta = lambda text: [roberta_sentiment(text)]
+
+            
+            # Initialize finbert_tone
+            finbert_tone_tokenizer = AutoTokenizer.from_pretrained("yiyanghkust/finbert-tone")
+            finbert_tone_model = AutoModelForSequenceClassification.from_pretrained(
+                "yiyanghkust/finbert-tone", 
+                use_safetensors=True,
+                device_map=device
+            )
+            
+            # Define custom pipeline instead of using transformers pipeline
+            def finbert_tone_sentiment(text):
+                inputs = finbert_tone_tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
+                outputs = finbert_tone_model(**inputs)
+                probs = outputs.logits.softmax(dim=1)[0]
+                pred_idx = probs.argmax().item()
+                labels = ["negative", "neutral", "positive"]
+                return {"label": labels[pred_idx], "score": probs[pred_idx].item()}
+                
+            self.finbert_tone = lambda text: [finbert_tone_sentiment(text)]
             
             # Initialize MT5 model with safetensors
             self.model_name = "google/mt5-small"