Update app.py

app.py

@@ -1,9 +1,12 @@
 import gradio as gr
-from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
 import json
 import torch
+from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
+from transformers import AutoModelForTokenClassification
+from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
 
-# Define hyperparameters
+
+# Define hyperparameters for QnA
 learning_rate = 3e-5
 batch_size = 16
 epochs = 3
@@ -12,12 +15,15 @@ warmup_steps = 100
 weight_decay = 0.01
 dropout_prob = 0.1
 gradient_clip_value = 1.0
-
-context_val = ''
-
 q_n_a_model_name = "deepset/roberta-base-squad2"
 q_n_a_model = AutoModelForQuestionAnswering.from_pretrained(q_n_a_model_name)
-tokenizer = AutoTokenizer.from_pretrained(q_n_a_model_name)  
+q_n_a_tokenizer = AutoTokenizer.from_pretrained(q_n_a_model_name)  
+
+
+classification_tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
+classification_model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased")
+
+
 
 context = gr.Textbox(label="Add the Context (Paragraph or texts) for which you want to get insights", lines=10, outputs="text")
 
@@ -29,7 +35,7 @@ def q_n_a_fn(context, text):
     q_n_a_model.to(device)
     
     # Convert inputs to tensors
-    inputs = tokenizer(QA_input["context"], QA_input["question"], return_tensors="pt", max_length=max_seq_length, truncation=True, padding="max_length").to(device)
+    inputs = q_n_a_tokenizer(QA_input["context"], QA_input["question"], return_tensors="pt", max_length=max_seq_length, truncation=True, padding="max_length").to(device)
 
     # Get predictions
     with torch.no_grad():
@@ -46,14 +52,21 @@ def q_n_a_fn(context, text):
     answer_tokens = inputs["input_ids"][0][start_idx : end_idx + 1]
 
     # Decode the answer tokens into a human-readable answer
-    answer = tokenizer.decode(inputs["input_ids"][0][start_idx:end_idx+1], skip_special_tokens=True)
+    answer = q_n_a_tokenizer.decode(inputs["input_ids"][0][start_idx:end_idx+1], skip_special_tokens=True)
 
     return answer
 
-def classification_fn(text):
-    return context
+def classification_fn(context):    
+    inputs = classification_tokenizer(context, return_tensors="pt")
+    with torch.no_grad():
+        logits = classification_model(**inputs).logits
+        class_probabilities = torch.softmax(logits, dim=1)
+
+    class_probabilities = torch.softmax(logits, dim=1)
+    class_probabilities = class_probabilities[0].tolist()  # Convert to a Python list
+    return {"POSITIVE": class_probabilities[0], "NEGATIVE": class_probabilities[1]}
 
-def translate_fn(text):
+def translate_fn(context, text):
     return context
     
 with gr.Blocks(theme='gradio/soft') as demo:
@@ -65,11 +78,11 @@ with gr.Blocks(theme='gradio/soft') as demo:
     
     with gr.Tab("Classifier"):
         with gr.Row():
-            gr.Interface(fn=classification_fn, inputs=[context], outputs="label")
+            gr.Interface(fn=classification_fn, inputs=[context], outputs=[gr.Label()])
     
     with gr.Tab("Translation"):
         with gr.Row():
-            gr.Interface(fn=translate_fn, inputs=[gr.Radio(["French", "Hindi", "Spanish"], label="Languages", info="Select language")], outputs="text")
+            gr.Interface(fn=translate_fn, inputs=[context, gr.Radio(["French", "Hindi", "Spanish"], label="Languages", info="Select language")], outputs="text")
 
     with gr.Tab("Summarization"):
         with gr.Row():