Update app.py

app.py

@@ -1,11 +1,67 @@
-from transformers import pipeline
+from transformers import BertForSequenceClassification, BertTokenizer
+import torch
 import gradio as gr
-pipeline = pipeline("text-classification" , model = "ProsusAI/finbert")
+tokenizer = BertTokenizer.from_pretrained('ProsusAI/finbert')
+model = BertForSequenceClassification.from_pretrained('ProsusAI/finbert')
+
+
 
 def predict(input_text):
-  predictions = pipeline(input_text )
-  return predictions[0]
+  tokens = tokenizer.encode_plus(input_text, add_special_tokens = False, return_tensors = 'pt')
+  input_id_chunks = tokens['input_ids'][0].split(510)
+  attention_mask_chunks = tokens['attention_mask'][0].split(510)
+  def get_input_ids_and_attention_mask_chunk():
+    """
+    This function splits the input_ids and attention_mask into chunks of size 'chunksize'. 
+    It also adds special tokens (101 for [CLS] and 102 for [SEP]) at the start and end of each chunk.
+    If the length of a chunk is less than 'chunksize', it pads the chunk with zeros at the end.
+    
+    Returns:
+        input_id_chunks (List[torch.Tensor]): List of chunked input_ids.
+        attention_mask_chunks (List[torch.Tensor]): List of chunked attention_masks.
+    """
+    chunksize = 512
+    input_id_chunks = list(tokens['input_ids'][0].split(chunksize - 2))
+    attention_mask_chunks = list(tokens['attention_mask'][0].split(chunksize - 2))
+    
+    for i in range(len(input_id_chunks)):
+        input_id_chunks[i] = torch.cat([
+            torch.tensor([101]), input_id_chunks[i], torch.tensor([102])
+        ])
+        
+        attention_mask_chunks[i] = torch.cat([
+            torch.tensor([1]), attention_mask_chunks[i], torch.tensor([1])
+        ])
+        
+        pad_length = chunksize - input_id_chunks[i].shape[0]
+        
+        if pad_length > 0:
+            input_id_chunks[i] = torch.cat([
+                input_id_chunks[i], torch.Tensor([0] * pad_length)
+            ])
+            attention_mask_chunks[i] = torch.cat([
+                attention_mask_chunks[i], torch.Tensor([0] * pad_length)
+            ])
+            
+    return input_id_chunks, attention_mask_chunks 
+  input_id_chunks, attention_mask_chunks = get_input_ids_and_attention_mask_chunk()
+  input_ids = torch.stack(input_id_chunks)
+  attention_mask = torch.stack(attention_mask_chunks)
+  input_dict = {
+      'input_ids' : input_ids.long(),
+      'attention_mask' : attention_mask.int()
+  }
+  outputs = model(**input_dict)
+  probabilities = torch.nn.functional.softmax(outputs[0], dim = -1 )
+  mean_probabilities = probabilities.mean(dim = 0)
+  output = torch.argmax(mean_probabilities).item()
+  if output==0:
+    return "positive"
+  elif output==1:
+    return "negative"
 
+  elif output==2 :
+    return "neutral"  
 
 gradio_app = gr.Interface(
   predict,