Views
No views yet
1docstring = "sentences of docstring"
2dome = DOME("dome_location")
3sentences, predictions = dome.predict(docstring)spacy
torch
transformers1"""
2Model is based on replication package for ICSE23 Paper Developer-Intent Driven Code Comment Generation.
3Initial solution: https://github.com/ICSE-DOME/DOME
4Pipeline consists of several parts:
5* split docstring into sentences
6* prepare input data for DOMEBertForClassification
7* predict class
8
9How to use:
10```python
11docstring = "sentences of docstring"
12dome = DOME("dome_location")
13sentences, predictions = dome.predict(docstring)
14```
15"""
16import re
17from typing import Tuple, List
18
19import spacy
20import torch
21import torch.nn as nn
22import torch.nn.functional as F
23from transformers import AutoTokenizer, RobertaConfig, RobertaModel
24
25MAX_LENGTH_BERT = 510
26
27
28class DOME:
29 """
30 End-to-end pipeline for docstring classification
31 * split sentences
32 * prepare inputs
33 * classify
34 """
35 def __init__(self, pretrained_model: str):
36 """
37 :param pretrained_model: location of pretrained model
38 """
39 self.model = DOMEBertForClassification.from_pretrained(pretrained_model)
40 self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model)
41 self.docstring2sentences = Docstring2Sentences()
42
43 def predict(self, docstring: str) -> Tuple[List[str], List[str]]:
44 """
45 Predict DOME classes for each sentence in docstring.
46 :param docstring: docstring to process
47 :return: tuple with list of sentences and list of predictions for each sentence.
48 """
49 sentences = self.docstring2sentences.docstring2sentences(docstring)
50 predictions = [self.model.predict(*dome_preprocess(tokenizer=self.tokenizer, comment=sentence))
51 for sentence in sentences]
52 return sentences, predictions
53
54
55class DOMEBertForClassification(RobertaModel):
56 """
57 A custom classification model based on the RobertaModel for intent classification.
58
59 This model extends the RobertaModel with additional linear layers to incorporate
60 comment length as an additional feature for classification tasks.
61 """
62
63 DOME_CLASS_NAMES = ["what", "why", "how-to-use", "how-it-is-done", "property", "others"]
64
65 def __init__(self, config: RobertaConfig):
66 """
67 Initialize the DOMEBertForClassification model.
68
69 :param config: The configuration information for the RobertaModel.
70 """
71 super().__init__(config)
72
73 # I omit possibility to configure number of classes and so on because we need to load pretrained model
74 # DOME layers for intent classification:
75 self.fc1 = nn.Linear(768 + 1, 768 // 3)
76 self.fc2 = nn.Linear(768 // 3, 6)
77 self.dropout = nn.Dropout(0.2)
78
79 def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor = None, comment_len: torch.Tensor = None) \
80 -> torch.Tensor:
81 """
82 Forward pass for the DOMEBertForClassification model.
83
84 :param input_ids: Tensor of token ids to be fed to a model.
85 :param attention_mask: Mask to avoid performing attention on padding token indices. Always equals 1.
86 :param comment_len: Tensor representing the length of comments. Equal 1 if comment has less than 3 words,
87 0 otherwise.
88 :return: The logits after passing through the model.
89 """
90 # Use the parent class's forward method to get the base outputs
91 outputs = super().forward(
92 input_ids=input_ids,
93 attention_mask=attention_mask
94 )
95 # Extract the pooled output (last hidden state of the [CLS] token)
96 pooled_output = outputs.pooler_output
97 # DOME custom layers:
98 comment_len = comment_len.view(-1, 1).float() # Ensure comment_len is correctly shaped
99 # DOME use comment len as additional feature
100 combined_input = torch.cat([pooled_output, comment_len], dim=-1)
101 x = self.dropout(F.relu(self.fc1(self.dropout(combined_input))))
102 logits = self.fc2(x)
103 return logits
104
105 def predict(self, input_ids: torch.Tensor, attention_mask: torch.Tensor = None, comment_len: torch.Tensor = None) \
106 -> str:
107 """
108 Predict class for tokenized docstring.
109
110 :param input_ids: Tensor of token ids to be fed to a model.
111 :param attention_mask: Mask to avoid performing attention on padding token indices. Always equals 1.
112 :param comment_len: Tensor representing the length of comments. Equal 1 if comment has less than 3 words,
113 0 otherwise.
114 :return: class
115 """
116 logits = self.forward(input_ids=input_ids, attention_mask=attention_mask, comment_len=comment_len)
117 return self.DOME_CLASS_NAMES[int(torch.argmax(logits, 1))]
118
119
120def dome_preprocess(tokenizer, comment):
121 """
122 DOME preprocessor - returns all required values for "DOMEBertForClassification.forward".
123 This function limits maximum number of tokens to fit into BERT.
124 :param tokenizer: tokenizer to use.
125 :param comment: text of sentence from docstring/comment that should be classified by DOMEBertForClassification.
126 :return: tuple with (input_ids, attention_mask, comment_len).
127 """
128 input_ids = tokenizer.convert_tokens_to_ids([tokenizer.cls_token] + tokenizer.tokenize(comment) +
129 [tokenizer.sep_token])[:MAX_LENGTH_BERT]
130 attention_mask = [1] * len(input_ids)
131 if len(comment.strip().split()) < 3:
132 comment_len = 1
133 else:
134 comment_len = 0
135 return (torch.tensor(input_ids).unsqueeze(0), torch.tensor(attention_mask).unsqueeze(0),
136 torch.tensor(comment_len).unsqueeze(0))
137
138
139class Docstring2Sentences:
140 """Helper class to split docstrings into sentences"""
141 def __init__(self):
142 self.spacy_nlp = spacy.load("en_core_web_sm")
143
144 @staticmethod
145 def split_docstring(docstring: str, delimiters: List[Tuple[str, str]]):
146 """
147 Splits the docstring into separate parts of text and code blocks, preserving the original formatting.
148
149 :param docstring: The docstring to split.
150 :param delimiters: A list of tuples, each containing start and end delimiters for code blocks.
151 :return: A list of strings, each either a text block or a code block.
152 """
153
154 # Escape delimiter parts for regex and create a combined pattern
155 escaped_delimiters = [tuple(map(re.escape, d)) for d in delimiters]
156 combined_pattern = '|'.join([f'({start}.*?{end})' for start, end in escaped_delimiters])
157
158 # Split using the combined pattern, preserving the delimiters
159 parts = re.split(combined_pattern, docstring, flags=re.DOTALL)
160
161 # Filter out empty strings
162 parts = [part for part in parts if part]
163
164 return parts
165
166 @staticmethod
167 def is_only_spaces_and_newlines(string):
168 """
169 Check if the given string contains only spaces and newlines.
170
171 :param string: The string to check.
172 :return: True if the string contains only spaces and newlines, False otherwise.
173 """
174 return bool(re.match(r'^[\s\n]+$', string))
175
176 def docstring2sentences(self, docstring):
177 """
178 Splits a docstring into individual sentences, preserving code blocks.
179
180 This function uses `docstring2parts` to split the docstring into parts based on predefined code block delimiters.
181 It then utilizes a SpaCy NLP model to split the non-code text parts into sentences.
182 Code blocks are kept intact as single elements.
183
184 :param docstring: The docstring to be processed, which may contain both regular text and code blocks.
185 :return: A list containing individual sentences and intact code blocks.
186 """
187 delimiters = [("@code", "@endcode"), ("\code", "\endcode")]
188 parts = self.split_docstring(docstring=docstring, delimiters=delimiters)
189 sentences = []
190 for part in parts:
191 if part[1:5] == "code" and part[-7:] == "endcode":
192 # code block
193 sentences.append(part)
194 else:
195 sentences.extend(sentence.text for sentence in self.spacy_nlp(part).sents)
196
197 return [sentence for sentence in sentences if not self.is_only_spaces_and_newlines(sentence)]
198