Update modeling_provence.py

modeling_provence.py

@@ -100,8 +100,8 @@ class Provence(DebertaV2PreTrainedModel):
     def process(
         self,
         question: Union[List[str], str],
-        context: Union[List[List[str]], List[str], str],
-        title: Optional[Union[List[List[str]], List[str], str]] = "first_sentence",
+        context: Union[List[List[str]], str],
+        title: Optional[Union[List[List[str]], str]] = "first_sentence",
         batch_size=32,
         threshold=0.1,
         always_select_title=False,
@@ -117,14 +117,10 @@ class Provence(DebertaV2PreTrainedModel):
             queries = question
         if type(context) == str:
             contexts = [[context]]
-        elif type(context) == list and type(context[0]) == str:
-            contexts = [context]
         else:
             contexts = context
         if type(title) == str:
             titles = [[title]]
-        elif type(title) == list and type(title[0]) == str:
-            titles = [title]
         else:
             titles = title
         assert (
@@ -235,13 +231,10 @@ class Provence(DebertaV2PreTrainedModel):
         if type(context) == str:
             selected_contexts = selected_contexts[0][0]
             reranking_scores = reranking_scores[0][0]
-        elif type(context) == list and type(context[0]) == str:
-            selected_contexts = selected_contexts[0]
-            reranking_scores = reranking_scores[0]
                 
         return {
-            "pruned_contexts": selected_contexts,
-            "reranking_scores": reranking_scores
+            "pruned_context": selected_contexts,
+            "reranking_score": reranking_scores
         }