Upload extract.py with huggingface_hub

extract.py

@@ -0,0 +1,86 @@
+from pydoc import text
+
+import torch
+from src.task.model import TinyTransformerLM, TinyClassifier, TinyAnsweringModel
+
+class TextExtractor:
+
+    def __init__(
+        self,
+        model: TinyAnsweringModel,
+        tokenizer,
+        context_length,
+        device=None,
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.context_length = context_length
+
+        if device is None:
+            device = (
+                "cuda"
+                if torch.cuda.is_available()
+                else "cpu"
+            )
+
+        self.device = torch.device(device)
+        self.model.to(self.device)
+        self.model.eval()
+
+    @torch.no_grad()
+    def extract(
+        self,
+        context_text: str,  #context containing the answer
+        question_text: str  #the actual question 
+        ):
+
+        input_text = (
+                self.tokenizer.CLS_TOKEN
+                + question_text
+                + self.tokenizer.SEP_TOKEN
+                + context_text
+                + self.tokenizer.SEP_TOKEN
+                )
+        input_ids = self.tokenizer.encode(input_text)
+
+        padded_ids, mask_ids = self.tokenizer.pad(
+            [input_ids],
+            max_length=self.context_length
+            )
+
+        token_tensor = torch.tensor(
+            padded_ids[0],
+            dtype=torch.long,
+            device=self.device,
+            )
+
+        padding_mask = torch.tensor(
+            mask_ids,
+            dtype=torch.bool,
+            device=self.device,
+            )
+
+        #Typically -> embedding -> transformer -> start,end logits
+        #Forward pass on the model is refusing :-(
+        start_logits, end_logits = self.model( 
+            tokens=token_tensor.unsqueeze(0),
+            padding_mask=padding_mask
+            )
+
+
+        # Remove batch dimension: [max_length]
+        start_logits = start_logits[0]
+        end_logits = end_logits[0]
+        
+        start_predictions = torch.argmax(
+            start_logits,
+            dim=-1,
+            )
+    
+        end_predictions = torch.argmax(
+            end_logits,
+            dim=-1,
+            )
+        
+        return token_tensor, start_predictions.item(), end_predictions.item()
+