zeguru commited on
Commit
f919114
·
verified ·
1 Parent(s): 9d6bce3

Delete files extract.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. extract.py +0 -86
extract.py DELETED
@@ -1,86 +0,0 @@
1
- from pydoc import text
2
-
3
- import torch
4
- from src.task.model import TinyTransformerLM, TinyClassifier, TinyAnsweringModel
5
-
6
- class TextExtractor:
7
-
8
- def __init__(
9
- self,
10
- model: TinyAnsweringModel,
11
- tokenizer,
12
- context_length,
13
- device=None,
14
- ):
15
- self.model = model
16
- self.tokenizer = tokenizer
17
- self.context_length = context_length
18
-
19
- if device is None:
20
- device = (
21
- "cuda"
22
- if torch.cuda.is_available()
23
- else "cpu"
24
- )
25
-
26
- self.device = torch.device(device)
27
- self.model.to(self.device)
28
- self.model.eval()
29
-
30
- @torch.no_grad()
31
- def extract(
32
- self,
33
- context_text: str, #context containing the answer
34
- question_text: str #the actual question
35
- ):
36
-
37
- input_text = (
38
- self.tokenizer.CLS_TOKEN
39
- + question_text
40
- + self.tokenizer.SEP_TOKEN
41
- + context_text
42
- + self.tokenizer.SEP_TOKEN
43
- )
44
- input_ids = self.tokenizer.encode(input_text)
45
-
46
- padded_ids, mask_ids = self.tokenizer.pad(
47
- [input_ids],
48
- max_length=self.context_length
49
- )
50
-
51
- token_tensor = torch.tensor(
52
- padded_ids[0],
53
- dtype=torch.long,
54
- device=self.device,
55
- )
56
-
57
- padding_mask = torch.tensor(
58
- mask_ids,
59
- dtype=torch.bool,
60
- device=self.device,
61
- )
62
-
63
- #Typically -> embedding -> transformer -> start,end logits
64
- #Forward pass on the model is refusing :-(
65
- start_logits, end_logits = self.model(
66
- tokens=token_tensor.unsqueeze(0),
67
- padding_mask=padding_mask
68
- )
69
-
70
-
71
- # Remove batch dimension: [max_length]
72
- start_logits = start_logits[0]
73
- end_logits = end_logits[0]
74
-
75
- start_predictions = torch.argmax(
76
- start_logits,
77
- dim=-1,
78
- )
79
-
80
- end_predictions = torch.argmax(
81
- end_logits,
82
- dim=-1,
83
- )
84
-
85
- return token_tensor, start_predictions.item(), end_predictions.item()
86
-