| |
| from typing import Any, Dict, List, Optional |
|
|
| import numpy as np |
|
|
| from .core import ResponsePreprocessor |
|
|
|
|
| class GroundingMixin: |
| """This class offers prompts to the grounding task""" |
| task_type: Optional[str] = None |
|
|
| _grounding_language_mixin = [0.8, 0.2] |
| _grounding_prompts = { |
| 'grounding': { |
| 'en': [('<ref-object>', '<bbox>'), ('The positions of <ref-object> is', '<bbox>'), |
| ('Find the positions of <ref-object>', '<bbox>'), ('Where is <ref-object>', '<bbox>'), |
| ('Find <ref-object>', '<bbox>'), ('Show me <ref-object>', '<bbox>'), |
| ('Detect <ref-object>', '<bbox>'), ('Locate <ref-object>', '<bbox>'), |
| ('Tell me the location of <ref-object>', '<bbox>'), ('Give the location of <ref-object>', '<bbox>'), |
| ('Provide the bounding box coordinate of <ref-object>', '<bbox>')], |
| 'zh': [('<ref-object>', '<bbox>'), ('<ref-object>的位置在图片中', '<bbox>'), ('<ref-object>在图片中', '<bbox>'), |
| ('<ref-object>在', '<bbox>'), ('找到<ref-object>的位置', '<bbox>'), ('<ref-object>在哪里', '<bbox>'), |
| ('提供<ref-object>的坐标位置', '<bbox>')] |
| }, |
| 'caption': { |
| 'en': [ |
| ('<bbox>', '<ref-object>'), |
| ('The object at position <bbox>', '<ref-object>'), |
| ('This <bbox> is', '<ref-object>'), |
| ('What is the object at <bbox>', '<ref-object>'), |
| ('Describe <bbox>', '<ref-object>'), |
| ('<bbox> is', '<ref-object>'), |
| ('The bounding box coordinate <bbox> contains', '<ref-object>'), |
| ], |
| 'zh': [ |
| ('<bbox>', '<ref-object>'), |
| ('<bbox>是什么', '<ref-object>'), |
| ('<bbox>的位置包含', '<ref-object>'), |
| ('描述<bbox>', '<ref-object>'), |
| ('<bbox>中是', '<ref-object>'), |
| ('坐标<bbox>描述了什么', '<ref-object>'), |
| ('描述<bbox>中的事物', '<ref-object>'), |
| ] |
| }, |
| } |
|
|
| def construct_grounding_prompt(self): |
| |
| lang = np.random.choice(['en', 'zh'], p=[0.8, 0.2]) |
| prompts = GroundingMixin._grounding_prompts[self.task_type][lang] |
| query, response = prompts[np.random.choice(range(len(prompts)))] |
| return query, response |
|
|
|
|
| class TextGenerationPreprocessor(ResponsePreprocessor): |
|
|
| def __init__(self, |
| *, |
| prompt: str, |
| query_tag: str = '{{QUERY}}', |
| columns: Optional[Dict[str, str]] = None, |
| **kwargs) -> None: |
| self.query_tag = query_tag |
| self.prompt = prompt |
| super().__init__(columns=columns, **kwargs) |
|
|
| def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: |
| row['query'] = self.prompt.replace(self.query_tag, row['query']) |
| return super().preprocess(row) |
|
|
|
|
| class ClsGenerationPreprocessor(ResponsePreprocessor): |
|
|
| def __init__(self, |
| labels: List[str], |
| *, |
| task: str, |
| is_pair_seq: bool = False, |
| columns: Optional[Dict[str, str]] = None, |
| **kwargs) -> None: |
| self.labels = labels |
| self.task = task |
| self.is_pair_seq = is_pair_seq |
|
|
| category = ', '.join(labels) |
| self.sentence2_key = 'sentence2' |
| self.label_key = 'label' |
| if is_pair_seq: |
| self.sentence_key = 'sentence1' |
| inputs = 'Sentence1: {sentence1}\nSentence2: {sentence2}' |
| else: |
| self.sentence_key = 'sentence' |
| inputs = 'Sentence: {sentence}' |
| self.prompt = f"""Task: {task} |
| {inputs} |
| Category: {category} |
| Output:""" |
| super().__init__(columns=columns, **kwargs) |
|
|
| def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: |
| label = row.pop(self.label_key, None) |
| if label is None: |
| return |
|
|
| if self.is_pair_seq: |
| query = self.prompt.format(sentence1=row.pop(self.sentence_key), sentence2=row.pop(self.sentence2_key)) |
| else: |
| query = self.prompt.format(sentence=row.pop(self.sentence_key)) |
| row['query'] = query |
| row['response'] = self.labels[int(label)] |
| return super().preprocess(row) |
|
|
