Upload HybridQAPipeline

hybrid_pipe.py

@@ -1,11 +1,24 @@
 from hybrid_model import HybridQAModel
-from transformers import QuestionAnsweringPipeline
+from transformers import QuestionAnsweringPipeline, PretrainedConfig
 
 class HybridQAPipeline(QuestionAnsweringPipeline):
     def __init__(self, model=None, tokenizer=None, **kwargs):
-        self.config = kwargs['custom']
+        extractive_id = "datarpit/distilbert-base-uncased-finetuned-natural-questions"
+        generative_id = "MaRiOrOsSi/t5-base-finetuned-question-answering"
+        self.config = HybridQAConfig(extractive_id, generative_id)
         super().__init__(model=model, tokenizer=tokenizer, **kwargs)
         self.model = HybridQAModel(self.config)
 
     def __call__(self, question, context):
-        return self.model.predict(question, context)
+        return self.model.predict(question, context)
+
+class HybridQAConfig(PretrainedConfig):
+    def __init__(
+        self, 
+        extractive_id=None,
+        generative_id = None,
+        **kwargs
+    ):
+        self.extractive_id = extractive_id
+        self.generative_id = generative_id
+        super().__init__(**kwargs)