chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,233 @@
|
||||
<h1 align="center">CodeR: Towards A Generalist Code Embedding Model</h1>
|
||||
<p align="center">
|
||||
<a href="https://huggingface.co/datasets/nebula2025/CodeR-Pile">
|
||||
<img alt="Build" src="https://img.shields.io/badge/🤗 Dataset-CodeR Pile-yellow">
|
||||
</a>
|
||||
<a href="https://huggingface.co/nebula2025/CodeR-full">
|
||||
<img alt="Build" src="https://img.shields.io/badge/🤗 Model-CodeR Full-green">
|
||||
</a>
|
||||
<a href="https://huggingface.co/nebula2025/CodeR-synthetic">
|
||||
<img alt="Build" src="https://img.shields.io/badge/🤗 Model-CodeR Synthetic-blue">
|
||||
</a>
|
||||
</p>
|
||||
|
||||
|
||||
This repo contains the data, training, and evaluation pipeline for CodeR / [BGE-Code-v1](https://huggingface.co/BAAI/bge-code-v1)
|
||||
|
||||
**[BGE-Code-v1](https://huggingface.co/BAAI/bge-code-v1)** is an LLM-based code embedding model that supports code retrieval, text retrieval, and multilingual retrieval. It primarily demonstrates the following capabilities:
|
||||
|
||||
- Superior Code Retrieval Performance: The model demonstrates exceptional code retrieval capabilities, supporting natural language queries in both English and Chinese, as well as 20 programming languages.
|
||||
- Robust Text Retrieval Capabilities: The model maintains strong text retrieval capabilities comparable to text embedding models of similar scale.
|
||||
- Extensive Multilingual Support: BGE-Code-v1 offers comprehensive multilingual retrieval capabilities, excelling in languages such as English, Chinese, Japanese, French, and more.
|
||||
|
||||
## :bell: News:
|
||||
|
||||
- 🥳 5/15/2025: We have released the CodeR! :fire:
|
||||
|
||||
## Usage
|
||||
|
||||
### Using FlagEmbedding
|
||||
|
||||
```
|
||||
git clone https://github.com/FlagOpen/FlagEmbedding.git
|
||||
cd FlagEmbedding
|
||||
pip install -e .
|
||||
from FlagEmbedding import FlagLLMModel
|
||||
queries = [
|
||||
"Delete the record with ID 4 from the 'Staff' table.",
|
||||
'Delete all records in the "Livestock" table where age is greater than 5'
|
||||
]
|
||||
documents = [
|
||||
"DELETE FROM Staff WHERE StaffID = 4;",
|
||||
"DELETE FROM Livestock WHERE age > 5;"
|
||||
]
|
||||
model = FlagLLMModel('BAAI/bge-code-v1',
|
||||
query_instruction_format="<instruct>{}\n<query>{}",
|
||||
query_instruction_for_retrieval="Given a question in text, retrieve SQL queries that are appropriate responses to the question.",
|
||||
trust_remote_code=True,
|
||||
use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation
|
||||
embeddings_1 = model.encode_queries(queries)
|
||||
embeddings_2 = model.encode_corpus(documents)
|
||||
similarity = embeddings_1 @ embeddings_2.T
|
||||
print(similarity)
|
||||
```
|
||||
|
||||
By default, FlagLLMModel will use all available GPUs when encoding. Please set `os.environ["CUDA_VISIBLE_DEVICES"]` to select specific GPUs. You also can set `os.environ["CUDA_VISIBLE_DEVICES"]=""` to make all GPUs unavailable.
|
||||
|
||||
### Using Sentence Transformers
|
||||
|
||||
```python
|
||||
from sentence_transformers import SentenceTransformer
|
||||
import torch
|
||||
|
||||
# Load the model, optionally in float16 precision for faster inference
|
||||
model = SentenceTransformer(
|
||||
"BAAI/bge-code-v1",
|
||||
trust_remote_code=True,
|
||||
model_kwargs={"torch_dtype": torch.float16},
|
||||
)
|
||||
|
||||
# Prepare a prompt given an instruction
|
||||
instruction = 'Given a question in text, retrieve SQL queries that are appropriate responses to the question.'
|
||||
prompt = f'<instruct>{instruction}\n<query>'
|
||||
# Prepare queries and documents
|
||||
queries = [
|
||||
"Delete the record with ID 4 from the 'Staff' table.",
|
||||
'Delete all records in the "Livestock" table where age is greater than 5'
|
||||
]
|
||||
documents = [
|
||||
"DELETE FROM Staff WHERE StaffID = 4;",
|
||||
"DELETE FROM Livestock WHERE age > 5;"
|
||||
]
|
||||
|
||||
# Compute the query and document embeddings
|
||||
query_embeddings = model.encode(queries, prompt=prompt)
|
||||
document_embeddings = model.encode(documents)
|
||||
|
||||
# Compute the cosine similarity between the query and document embeddings
|
||||
similarities = model.similarity(query_embeddings, document_embeddings)
|
||||
print(similarities)
|
||||
```
|
||||
|
||||
### Using HuggingFace Transformers
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from torch import Tensor
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
|
||||
|
||||
def last_token_pool(last_hidden_states: Tensor,
|
||||
attention_mask: Tensor) -> Tensor:
|
||||
left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
|
||||
if left_padding:
|
||||
return last_hidden_states[:, -1]
|
||||
else:
|
||||
sequence_lengths = attention_mask.sum(dim=1) - 1
|
||||
batch_size = last_hidden_states.shape[0]
|
||||
return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
|
||||
|
||||
|
||||
def get_detailed_instruct(task_description: str, query: str) -> str:
|
||||
return f'<instruct>{task_description}\n<query>{query}'
|
||||
|
||||
|
||||
instruction = 'Given a question in text, retrieve SQL queries that are appropriate responses to the question.'
|
||||
queries = [
|
||||
"Delete the record with ID 4 from the 'Staff' table.",
|
||||
'Delete all records in the "Livestock" table where age is greater than 5'
|
||||
]
|
||||
documents = [
|
||||
"DELETE FROM Staff WHERE StaffID = 4;",
|
||||
"DELETE FROM Livestock WHERE age > 5;"
|
||||
]
|
||||
input_texts = queries + documents
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-code-v1', trust_remote_code=True)
|
||||
model = AutoModel.from_pretrained('BAAI/bge-code-v1', trust_remote_code=True)
|
||||
model.eval()
|
||||
|
||||
max_length = 4096
|
||||
# Tokenize the input texts
|
||||
batch_dict = tokenizer(input_texts, max_length=max_length, padding=True, truncation=True, return_tensors='pt', pad_to_multiple_of=8)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = model(**batch_dict)
|
||||
embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
|
||||
|
||||
# normalize embeddings
|
||||
embeddings = F.normalize(embeddings, p=2, dim=1)
|
||||
scores = (embeddings[:2] @ embeddings[2:].T) * 100
|
||||
print(scores.tolist())
|
||||
```
|
||||
|
||||
## Evaluation
|
||||
|
||||
**BGE-Code-v1** achieves state-of-the-art performance on both the CoIR and CodeRAG benchmarks.
|
||||
|
||||
- CoIR
|
||||
|
||||
| | CodeXEmbed-2B | CodeXEmbed-7B | Voyage-Code-002 | Voyage-Code-003 | BGE-Code-v1 |
|
||||
| --------------------- | ------------- | ------------- | --------------- | --------------- | ----------- |
|
||||
| **Apps** | 76.86 | 85.38 | 26.52 | 93.62 | 98.08 |
|
||||
| **CosQA** | 40.47 | 42.47 | 29.79 | 34.45 | 46.72 |
|
||||
| **Text2SQL** | 78.42 | 78.94 | 69.26 | 62.87 | 64.35 |
|
||||
| **CSN** | 87.87 | 89.67 | 81.79 | 89.35 | 89.53 |
|
||||
| **CSN-CCR** | 97.66 | 97.95 | 73.45 | 90.05 | 98.30 |
|
||||
| **CodeTrans-Contest** | 90.30 | 94.45 | 72.77 | 94.96 | 94.38 |
|
||||
| **CodeTrans-DL** | 38.57 | 40.46 | 27.48 | 38.57 | 46.13 |
|
||||
| **StackOverFlow-QA** | 94.47 | 96.33 | 67.68 | 97.17 | 95.35 |
|
||||
| **CodeFeedBack-ST** | 86.36 | 87.53 | 65.35 | 90.67 | 90.56 |
|
||||
| **CodeFeedBack-MT** | 65.51 | 68.83 | 28.74 | 93.58 | 94.38 |
|
||||
| **AVG** | **75.65** | **78.20** | **56.26** | **78.53** | **81.77** |
|
||||
|
||||
- CodedRAG
|
||||
|
||||
| | HummanEval | MBPP | DS-1000 | ODEX | RepoEval | SWE-bench-Lite | AVG |
|
||||
| --------------- | ---------- | ---- | ------- | ---- | -------- | -------------- | -------- |
|
||||
| SFR | 100.0 | 99.0 | 19.3 | 37.1 | 83.8 | 62.7 | **67.0** |
|
||||
| Jina-v2-code | 100.0 | 97.7 | 26.2 | 19.9 | 90.5 | 58.3 | **65.4** |
|
||||
| CodeXEmbed-2B | 100.0 | 97.4 | 25.4 | 23.9 | 88.7 | 52.4 | **64.6** |
|
||||
| Voyage-Code-002 | 100.0 | 99.0 | 33.1 | 26.6 | 94.3 | 29.1 | **63.7** |
|
||||
| BGE-Code-v1 | 100.0 | 99.2 | 40.9 | 36.1 | 93.1 | 67.4 | **72.8** |
|
||||
|
||||
### Instructions for Evaluation
|
||||
|
||||
```python
|
||||
{
|
||||
"Apps": "Given a code contest problem description, retrieve relevant code that can help solve the problem.",
|
||||
"CosQA": "Given a web search query, retrieve relevant code that can help answer the query.",
|
||||
"Text2SQL": "Given a question in text, retrieve SQL queries that are appropriate responses to the question.",
|
||||
"CSN": "Given a piece of code, retrieve the document string that summarizes the code.",
|
||||
"CSN-CCR": "Given a piece of code segment, retrieve the code segment that is the latter part of the code.",
|
||||
"CodeTrans-DL": "Given a piece of code, retrieve code that is semantically equivalent to the input code.",
|
||||
"CodeTrans-Contest": "Given a piece of Python code, retrieve C++ code that is semantically equivalent to the input code.",
|
||||
"StackOverFlow-QA": "Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.",
|
||||
"CodeFeedBack-ST": "Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.",
|
||||
"CodeFeedBack-MT": "Given a multi-turn conversation history that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.",
|
||||
"HummanEval": "Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.",
|
||||
"MBPP": "Given a textual explanation of code functionality, retrieve the corresponding code implementation.",
|
||||
"DS-1000": "Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.",
|
||||
"ODEX": "Given a question, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.",
|
||||
"RepoEval": "Given a piece of code segment, retrieve the code segment that is the latter part of the code.",
|
||||
"SWE-bench-Lite": "Given a code snippet containing a bug and a natural language description of the bug or error, retrieve code snippets that demonstrate solutions or fixes for similar bugs or errors (the desired documents)."
|
||||
}
|
||||
```
|
||||
|
||||
### Evaluation script
|
||||
|
||||
#### CoIR
|
||||
|
||||
For CoIR, we use the [CoIR](https://github.com/CoIR-team/coir) evaluation script:
|
||||
|
||||
```shell
|
||||
cd ./evaluation/coir_eval
|
||||
### clone coir
|
||||
mkdir test
|
||||
cd ./test
|
||||
git clone https://github.com/CoIR-team/coir.git
|
||||
mv ./coir/coir ../
|
||||
cd ..
|
||||
rm -rf ./test
|
||||
### evaluate
|
||||
bash eval.sh
|
||||
```
|
||||
|
||||
### CodeRAG
|
||||
|
||||
For CodeRAG, we use the [CodeRAG](https://github.com/code-rag-bench/code-rag-bench) evaluation script:
|
||||
|
||||
```shell
|
||||
cd ./evaluation/coderag_eval
|
||||
### clone coderag
|
||||
git clone https://github.com/code-rag-bench/code-rag-bench.git
|
||||
## You need prepare environment according to README.md
|
||||
rm -rf ./code-rag-bench/retrieval/create
|
||||
cp -r ./test/* ./code-rag-bench/retrieval/
|
||||
### prepare data
|
||||
bash prepare_data.sh
|
||||
### evaluate
|
||||
bash eval.sh
|
||||
```
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,104 @@
|
||||
import os
|
||||
import random
|
||||
import datasets
|
||||
from tqdm import tqdm
|
||||
from typing import List, Tuple
|
||||
|
||||
from utils import clean_code
|
||||
from constant import DocLength
|
||||
|
||||
|
||||
class CorpusGenerator:
|
||||
def __init__(
|
||||
self,
|
||||
cache_dir: str = None,
|
||||
):
|
||||
self.cache_dir = cache_dir
|
||||
|
||||
def _load_corpus(self, corpus_dir: str, doc_length: List[str], external_path: List[str],
|
||||
source_language: str, stop_threshold: int = -1):
|
||||
"""
|
||||
Load availavle documents for a given task from the CoIR-Retrieval dataset.
|
||||
"""
|
||||
|
||||
corpus_list = []
|
||||
|
||||
if corpus_dir is not None and os.path.exists(corpus_dir):
|
||||
file_list = os.listdir(corpus_dir)
|
||||
random.shuffle(file_list)
|
||||
|
||||
for file in file_list:
|
||||
flag = False
|
||||
if not file.endswith('.jsonl'):
|
||||
flag = False
|
||||
for d_length in doc_length:
|
||||
d_length = DocLength[d_length].value
|
||||
if d_length in file:
|
||||
flag = True
|
||||
if flag is False:
|
||||
continue
|
||||
file_path = os.path.join(corpus_dir, file)
|
||||
corpus = datasets.load_dataset('json', data_files=file_path, cache_dir=self.cache_dir)['train']
|
||||
for data in tqdm(corpus, desc="Loading corpus"):
|
||||
if source_language is None:
|
||||
lang = os.path.basename(corpus_dir)
|
||||
data['language'] = lang
|
||||
else:
|
||||
data['language'] = source_language
|
||||
|
||||
text = clean_code(data["text"], data["language"], length_threshold=200)
|
||||
data["text"] = text
|
||||
if text != '':
|
||||
corpus_list.append(data)
|
||||
|
||||
if stop_threshold > 0 and len(corpus_list) > stop_threshold:
|
||||
break
|
||||
break
|
||||
|
||||
for ep in external_path:
|
||||
if os.path.exists(ep):
|
||||
corpus = datasets.load_dataset('json', data_files=ep, cache_dir=self.cache_dir)['train']
|
||||
for data in tqdm(corpus, desc="Loading corpus"):
|
||||
if source_language is None:
|
||||
lang = os.path.basename(os.path.dirname(ep))
|
||||
data['language'] = lang
|
||||
else:
|
||||
data['language'] = source_language
|
||||
|
||||
# useful when the text is not present in the data
|
||||
if "text" not in data:
|
||||
data["text"] = data["pos"][0]
|
||||
|
||||
corpus_list.append(data)
|
||||
text = clean_code(data["text"], lang, length_threshold=200)
|
||||
data["text"] = text
|
||||
if text != '':
|
||||
corpus_list.append(data)
|
||||
|
||||
return corpus_list
|
||||
|
||||
def run(
|
||||
self,
|
||||
num_samples: int = -1,
|
||||
max_corpus: int = -1,
|
||||
corpus_dir: str = None,
|
||||
doc_length: List[str] = ["len_0_500"],
|
||||
external_path: List[str] = None,
|
||||
source_language: str = None
|
||||
) -> Tuple[List[dict], List[dict]]:
|
||||
stop_threshold = max(num_samples * 10, max_corpus * 2)
|
||||
corpus_list = self._load_corpus(
|
||||
corpus_dir, doc_length, external_path, source_language, stop_threshold
|
||||
)
|
||||
|
||||
if num_samples > 0 and num_samples < len(corpus_list):
|
||||
small_corpus_list = random.sample(corpus_list, num_samples)
|
||||
else:
|
||||
small_corpus_list = corpus_list
|
||||
|
||||
if max_corpus > 0 and max_corpus < len(corpus_list):
|
||||
corpus_list = random.sample(corpus_list, max_corpus)
|
||||
else:
|
||||
corpus_list = corpus_list
|
||||
|
||||
return small_corpus_list, corpus_list
|
||||
@@ -0,0 +1,127 @@
|
||||
import os
|
||||
import json
|
||||
from constant import Language, CodeLanguage, TaskType, CODE_TRANSLATION_RETRIEVAL_PAIRS, \
|
||||
get_pos_as_input_by_task_type
|
||||
|
||||
|
||||
def format_generated_examples(
|
||||
file_path: str,
|
||||
save_path: str,
|
||||
task_type: TaskType
|
||||
):
|
||||
if os.path.exists(save_path):
|
||||
return
|
||||
|
||||
if not os.path.exists(file_path):
|
||||
print("====================================")
|
||||
print("Warning: file not found! Maybe need to generate it first.")
|
||||
print(f"file_path: {file_path}")
|
||||
return
|
||||
|
||||
pos_as_input = get_pos_as_input_by_task_type(task_type)
|
||||
|
||||
data_list = []
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
for line in f.readlines():
|
||||
data = json.loads(line)
|
||||
|
||||
if pos_as_input:
|
||||
_input = data["pos"][0]
|
||||
_output = data["query"]
|
||||
else:
|
||||
_input = data["query"]
|
||||
_output = data["pos"][0]
|
||||
|
||||
if 'provided' in _input:
|
||||
continue
|
||||
if len(_input) > 12000 or len(_output) > 12000:
|
||||
continue
|
||||
|
||||
data_list.append({
|
||||
"input": _input,
|
||||
"output": _output
|
||||
})
|
||||
|
||||
if len(data_list) == 0:
|
||||
print("====================================")
|
||||
print("Warning: no data found!")
|
||||
print(f"file_path: {file_path}")
|
||||
return
|
||||
|
||||
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
||||
with open(save_path, "w", encoding="utf-8") as f:
|
||||
json.dump(data_list, f, indent=4, ensure_ascii=False)
|
||||
|
||||
|
||||
def main():
|
||||
original_gen_examples_dir = "./examples"
|
||||
|
||||
formatted_examples_dir = "./filtered_for_generation"
|
||||
|
||||
for language in Language:
|
||||
for task_type in TaskType:
|
||||
if task_type == TaskType.code_translation_retrieval:
|
||||
for code_language_pair in CODE_TRANSLATION_RETRIEVAL_PAIRS:
|
||||
code_language, tgt_code_language = code_language_pair
|
||||
|
||||
file_path = os.path.join(
|
||||
original_gen_examples_dir,
|
||||
language.name, task_type.name, f"{language.name}-{code_language.name}-to-{tgt_code_language.name}-triplets.jsonl"
|
||||
)
|
||||
save_path = os.path.join(
|
||||
formatted_examples_dir,
|
||||
language.name, task_type.name, f"{code_language.name}-to-{tgt_code_language.name}_sample_examples.json"
|
||||
)
|
||||
|
||||
format_generated_examples(file_path, save_path, task_type)
|
||||
|
||||
for code_language_pair in CODE_TRANSLATION_RETRIEVAL_PAIRS:
|
||||
tgt_code_language, code_language = code_language_pair
|
||||
|
||||
file_path = os.path.join(
|
||||
original_gen_examples_dir,
|
||||
language.name, task_type.name, f"{language.name}-{code_language.name}-to-{tgt_code_language.name}-triplets.jsonl"
|
||||
)
|
||||
save_path = os.path.join(
|
||||
formatted_examples_dir,
|
||||
language.name, task_type.name, f"{code_language.name}-to-{tgt_code_language.name}_sample_examples.json"
|
||||
)
|
||||
|
||||
format_generated_examples(file_path, save_path, task_type)
|
||||
|
||||
elif task_type == TaskType.text2sql_retrieval:
|
||||
file_path = os.path.join(
|
||||
original_gen_examples_dir,
|
||||
language.name, task_type.name, f"{language.name}-sql-triplets.jsonl"
|
||||
)
|
||||
save_path = os.path.join(
|
||||
formatted_examples_dir,
|
||||
language.name, task_type.name, "sql_sample_examples.json"
|
||||
)
|
||||
|
||||
format_generated_examples(file_path, save_path, task_type)
|
||||
|
||||
elif task_type == TaskType.code_context_retrieval:
|
||||
continue
|
||||
|
||||
else:
|
||||
for code_language in CodeLanguage:
|
||||
if code_language == CodeLanguage.null:
|
||||
continue
|
||||
|
||||
file_path = os.path.join(
|
||||
original_gen_examples_dir,
|
||||
language.name, task_type.name, f"{language.name}-{code_language.name}-triplets.jsonl"
|
||||
)
|
||||
save_path = os.path.join(
|
||||
formatted_examples_dir,
|
||||
language.name, task_type.name, f"{code_language.name}_sample_examples.json"
|
||||
)
|
||||
|
||||
format_generated_examples(file_path, save_path, task_type)
|
||||
|
||||
print("All done!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,134 @@
|
||||
import os
|
||||
import time
|
||||
import openai
|
||||
import random
|
||||
import tiktoken
|
||||
import threading
|
||||
from openai import OpenAI, AzureOpenAI
|
||||
from typing import Tuple
|
||||
|
||||
|
||||
class LLM:
|
||||
def __init__(
|
||||
self,
|
||||
model: str="Qwen2-5-Coder-32B-Instruct",
|
||||
model_type: str = "open-source",
|
||||
port: int = 8000,
|
||||
):
|
||||
if model_type == "open-source":
|
||||
self.client = OpenAI(
|
||||
api_key="EMPTY",
|
||||
base_url=f"http://localhost:{port}/v1/"
|
||||
)
|
||||
elif model_type == "azure":
|
||||
self.client = AzureOpenAI(
|
||||
api_key=os.getenv("OPENAI_API_KEY"),
|
||||
api_version=os.getenv("AZURE_API_VERSION", "2024-02-01"),
|
||||
azure_endpoint=os.getenv("AZURE_ENDPOINT"),
|
||||
azure_deployment=os.getenv("OPENAI_DEPLOYMENT_NAME", 'gpt-35-turbo')
|
||||
)
|
||||
elif model_type == "openai":
|
||||
self.client = OpenAI(
|
||||
api_key=os.getenv("OPENAI_API_KEY"),
|
||||
base_url=os.getenv("OPENAI_BASE_URL", None)
|
||||
)
|
||||
else:
|
||||
raise ValueError("model_type must be one of ['open-source', 'azure', 'openai']")
|
||||
|
||||
self.model = model
|
||||
self.tokenizer = tiktoken.get_encoding("o200k_base")
|
||||
|
||||
def split_text(self, text: str, anchor_points: Tuple[float, float] = (0.4, 0.7)):
|
||||
token_ids = self.tokenizer.encode(text)
|
||||
anchor_point = random.uniform(anchor_points[0], anchor_points[1])
|
||||
split_index = int(len(token_ids) * anchor_point)
|
||||
return self.tokenizer.decode(token_ids[:split_index]), self.tokenizer.decode(token_ids[split_index:])
|
||||
|
||||
def chat(
|
||||
self,
|
||||
prompt: str,
|
||||
max_tokens: int = 8192,
|
||||
logit_bais: dict = None,
|
||||
n: int = 1,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 0.6,
|
||||
repetition_penalty: float = 1.0,
|
||||
remove_thinking: bool = True,
|
||||
timeout: int = 90,
|
||||
):
|
||||
endure_time = 0
|
||||
endure_time_limit = timeout * 2
|
||||
|
||||
def create_completion(results):
|
||||
try:
|
||||
completion = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
max_tokens=max_tokens,
|
||||
logit_bias=logit_bais if logit_bais is not None else {},
|
||||
n=n,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
extra_body={'repetition_penalty': repetition_penalty},
|
||||
timeout=timeout,
|
||||
)
|
||||
results["content"] = [x.message.content for x in completion.choices[:n]]
|
||||
except openai.BadRequestError as e:
|
||||
# The response was filtered due to the prompt triggering Azure OpenAI's content management policy.
|
||||
results["content"] = [None for _ in range(n)]
|
||||
except openai.APIConnectionError as e:
|
||||
results["error"] = f'APIConnectionError({e})'
|
||||
except openai.RateLimitError as e:
|
||||
results["error"] = f'RateLimitError({e})'
|
||||
except Exception as e:
|
||||
results["error"] = f"Error: {e}"
|
||||
|
||||
while True:
|
||||
results = {"content": None, "error": None}
|
||||
completion_thread = threading.Thread(target=create_completion, args=(results,))
|
||||
completion_thread.start()
|
||||
|
||||
start_time = time.time()
|
||||
while completion_thread.is_alive():
|
||||
elapsed_time = time.time() - start_time
|
||||
if elapsed_time > endure_time_limit:
|
||||
print("Completion timeout exceeded. Aborting...")
|
||||
return [None for _ in range(n)]
|
||||
time.sleep(1)
|
||||
|
||||
# If an error occurred during result processing
|
||||
if results["error"]:
|
||||
if endure_time >= endure_time_limit:
|
||||
print(f'{results["error"]} - Skip this prompt.')
|
||||
return [None for _ in range(n)]
|
||||
print(f"{results['error']} - Waiting for 5 seconds...")
|
||||
endure_time += 5
|
||||
time.sleep(5)
|
||||
continue
|
||||
|
||||
content_list = results["content"]
|
||||
if remove_thinking:
|
||||
content_list = [x.split('</think>')[-1].strip('\n').strip() if x is not None else None for x in content_list]
|
||||
return content_list
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
llm = LLM(
|
||||
model="gpt-4o-mini-2024-07-18",
|
||||
model_type="openai"
|
||||
)
|
||||
|
||||
prompt = "hello, who are you?"
|
||||
response = llm.chat(prompt)[0]
|
||||
print(response)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
llm = LLM(
|
||||
model="gpt-4o-mini-2024-07-18",
|
||||
model_type="openai"
|
||||
)
|
||||
|
||||
prompt = "hello, who are you?"
|
||||
response = llm.chat(prompt)[0]
|
||||
print(response)
|
||||
@@ -0,0 +1,368 @@
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
import gc
|
||||
import torch
|
||||
import argparse
|
||||
import random
|
||||
from hashlib import md5
|
||||
import multiprocessing as mp
|
||||
from typing import List, Optional
|
||||
|
||||
from constant import TaskType, Language, CodeLanguage, NUM_HARD_NEGATIVES
|
||||
from corpus_generator import CorpusGenerator
|
||||
from triplet_generator import TripletGenerator
|
||||
from search import get_top1
|
||||
|
||||
|
||||
def compute_md5(text: str):
|
||||
return md5(text.encode()).hexdigest()
|
||||
|
||||
|
||||
def get_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
'--task_type',
|
||||
type=str,
|
||||
required=True,
|
||||
help='The task type to generate data for',
|
||||
choices=[t.name for t in TaskType]
|
||||
)
|
||||
parser.add_argument(
|
||||
'--code_language',
|
||||
type=str,
|
||||
required=True,
|
||||
help='The code language to generate questions for.',
|
||||
choices=[c.name for c in CodeLanguage]
|
||||
)
|
||||
parser.add_argument(
|
||||
'--corpus_root',
|
||||
type=str,
|
||||
required=True,
|
||||
help='The root directory of the corpus data.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--save_dir',
|
||||
type=str,
|
||||
required=True,
|
||||
help='The path to save the generated data'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--examples_dir',
|
||||
type=str,
|
||||
default=None,
|
||||
help='The path to the examples directory. If not None, the examples will be used for few-shot generation.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--num_examples',
|
||||
type=int,
|
||||
default=3,
|
||||
help='The number of examples to use for few-shot generation. Default: 3'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--cache_dir',
|
||||
type=str,
|
||||
default=None,
|
||||
help='The cache directory'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--language',
|
||||
type=str,
|
||||
default='en',
|
||||
help='The language to generate for. ISO 639-1 code. Default: en',
|
||||
choices=[l.name for l in Language]
|
||||
)
|
||||
parser.add_argument(
|
||||
'--tgt_code_language',
|
||||
type=str,
|
||||
default=None,
|
||||
help='The target code language to generate code translations for.',
|
||||
choices=[c.name for c in CodeLanguage]
|
||||
)
|
||||
parser.add_argument(
|
||||
'--num_samples',
|
||||
type=int,
|
||||
default=-1,
|
||||
help='The number of examples to use for generation. Default: -1. Use all available examples.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--model',
|
||||
type=str,
|
||||
default='Qwen2.5-72B-Instruct',
|
||||
help='The model to use for generation. Default: Qwen2.5-72B-Instruct'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--model_type',
|
||||
type=str,
|
||||
default='open-source',
|
||||
help='The type of model to use for generation. Default: open-source',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--port',
|
||||
type=int,
|
||||
default=8000,
|
||||
help='The port for vllm.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--num_processes',
|
||||
type=int,
|
||||
default=1,
|
||||
help='The number of processes to use for generation. Default: 1'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--doc_length',
|
||||
type=str,
|
||||
default='len_0_500',
|
||||
help='The corpus length used to load dataset. Default: len_0_500'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--external_path',
|
||||
type=str,
|
||||
default='',
|
||||
help='The corpus length used to load dataset. Default: len_0_500'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--sim_model_name',
|
||||
type=str,
|
||||
default=None,
|
||||
help='The language of source corpus.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--max_corpus',
|
||||
type=int,
|
||||
default=500000,
|
||||
help='The max num of corpus to load.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--overwrite',
|
||||
action='store_true',
|
||||
help='Whether to overwrite the existing data.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--debug_mode',
|
||||
action='store_true',
|
||||
help='Whether to open debug mode.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--gen_hard_neg',
|
||||
action='store_true',
|
||||
help='Whether to generate hard negatives.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--seed',
|
||||
type=int,
|
||||
default=None,
|
||||
help='Random seed for generating triplets using the same positive. Default: 42'
|
||||
)
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def gen_triplets(
|
||||
model: str,
|
||||
model_type: str,
|
||||
port: int,
|
||||
positives: List[dict],
|
||||
task_type: str,
|
||||
language: str,
|
||||
code_language: str,
|
||||
tgt_code_language: str,
|
||||
examples_pool: Optional[List[dict]] = None,
|
||||
num_examples: int = 3,
|
||||
tqdm_desc: str = "Generating triplets",
|
||||
thread_count: int = 1,
|
||||
gen_cache_dir: Optional[str] = None,
|
||||
debug_mode: bool = False,
|
||||
gen_hard_neg: bool = False,
|
||||
):
|
||||
triplet_generator = TripletGenerator(model, model_type, port, cache_dir=gen_cache_dir)
|
||||
triplets = triplet_generator.run(
|
||||
positives=positives,
|
||||
task_type=task_type,
|
||||
language=language,
|
||||
code_language=code_language,
|
||||
tgt_code_language=tgt_code_language,
|
||||
examples_pool=examples_pool,
|
||||
num_examples=num_examples,
|
||||
tqdm_desc=tqdm_desc,
|
||||
thread_count=thread_count,
|
||||
debug_mode=debug_mode,
|
||||
gen_hard_neg=gen_hard_neg,
|
||||
num_negatives=NUM_HARD_NEGATIVES,
|
||||
)
|
||||
return triplets
|
||||
|
||||
|
||||
def get_save_path(
|
||||
save_dir: str,
|
||||
task_type: str,
|
||||
language: str,
|
||||
code_language: str,
|
||||
tgt_code_language: Optional[str] = None
|
||||
):
|
||||
save_dir = os.path.join(save_dir, language, task_type)
|
||||
if tgt_code_language is not None:
|
||||
file_name = f"{language}-{code_language}-to-{tgt_code_language}-triplets.jsonl"
|
||||
else:
|
||||
file_name = f"{language}-{code_language}-triplets.jsonl"
|
||||
save_path = os.path.join(save_dir, file_name)
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
return save_path
|
||||
|
||||
|
||||
def save_triplets(
|
||||
triplets: list,
|
||||
save_dir: str,
|
||||
task_type: str,
|
||||
language: str,
|
||||
code_language: str,
|
||||
tgt_code_language: Optional[str] = None
|
||||
):
|
||||
if len(triplets) == 0:
|
||||
print(f"No triplets to save: {task_type} | {language} | {code_language} | {tgt_code_language}")
|
||||
return
|
||||
|
||||
save_path = get_save_path(save_dir, task_type, language, code_language, tgt_code_language)
|
||||
query_md5s = set()
|
||||
pos_md5s = set()
|
||||
old_triplets = []
|
||||
if os.path.exists(save_path):
|
||||
with open(save_path, "r", encoding="utf-8") as f:
|
||||
for line in f.readlines():
|
||||
triplet = json.loads(line)
|
||||
old_triplets.append(triplet)
|
||||
query_md5s.add(compute_md5(triplet['query']))
|
||||
pos_md5s.add(compute_md5(triplet['pos'][0]))
|
||||
|
||||
with open(save_path, 'w', encoding='utf-8') as f:
|
||||
for triplet in old_triplets:
|
||||
f.write(json.dumps(triplet, ensure_ascii=False) + '\n')
|
||||
|
||||
for triplet in triplets:
|
||||
_query_md5 = compute_md5(triplet['query'])
|
||||
_pos_md5 = compute_md5(triplet['pos'][0])
|
||||
if _query_md5 in query_md5s or _pos_md5 in pos_md5s:
|
||||
continue
|
||||
f.write(json.dumps(triplet, ensure_ascii=False) + '\n')
|
||||
print(f"Triplets saved to {save_path}")
|
||||
|
||||
|
||||
def main(args):
|
||||
# set seed
|
||||
seed = args.seed
|
||||
if seed is not None:
|
||||
print(f"------------------- Seed set to {seed} -------------------")
|
||||
random.seed(seed)
|
||||
|
||||
model = args.model
|
||||
model_type = args.model_type
|
||||
port = args.port
|
||||
|
||||
num_samples = args.num_samples
|
||||
|
||||
task_type = args.task_type
|
||||
language = args.language
|
||||
code_language = args.code_language
|
||||
tgt_code_language = args.tgt_code_language
|
||||
|
||||
corpus_root = args.corpus_root
|
||||
corpus_dir = os.path.join(corpus_root, code_language)
|
||||
doc_length = args.doc_length.split()
|
||||
external_path = args.external_path.split()
|
||||
|
||||
save_dir = args.save_dir
|
||||
cache_dir = args.cache_dir
|
||||
num_processes = min(args.num_processes, int(mp.cpu_count() * 0.8))
|
||||
overwrite = args.overwrite
|
||||
debug_mode = args.debug_mode
|
||||
gen_hard_neg = args.gen_hard_neg
|
||||
|
||||
save_path = get_save_path(save_dir, task_type, language, code_language, tgt_code_language)
|
||||
# if os.path.exists(save_path) and not overwrite:
|
||||
# data = []
|
||||
# with open(save_path) as f:
|
||||
# for line in f:
|
||||
# data.append(json.loads(line))
|
||||
# if len(data) >= num_samples * 0.8:
|
||||
# print(f"Triplets already exist at {save_path}. Skipping generation.")
|
||||
# return
|
||||
# else:
|
||||
# print(f"Triplets already exist at {save_path}. But samples is really small, continue generation.")
|
||||
# num_samples = int((num_samples - len(data)) * 1.25) # consider the filtered samples
|
||||
|
||||
corpus_generator = CorpusGenerator(cache_dir)
|
||||
|
||||
examples_dir = args.examples_dir
|
||||
num_examples = args.num_examples
|
||||
if examples_dir is not None:
|
||||
# if task_type in ["single_turn_code_qa", "multi_turn_code_qa"]:
|
||||
# examples_path = os.path.join(examples_dir, language, task_type, "sample_examples.json")
|
||||
if task_type in ["code_translation_retrieval"]:
|
||||
examples_path = os.path.join(examples_dir, language, task_type,
|
||||
f"{code_language}-to-{tgt_code_language}_sample_examples.json")
|
||||
else:
|
||||
examples_path = os.path.join(examples_dir, language, task_type, f"{code_language}_sample_examples.json")
|
||||
try:
|
||||
with open(examples_path, 'r', encoding='utf-8') as f:
|
||||
examples_pool = json.load(f)
|
||||
examples_pool = random.sample(examples_pool,
|
||||
min(30, len(examples_pool))) # sample 30 examples for few-shot generation
|
||||
except:
|
||||
print(f'Error for loading examples from {examples_path}')
|
||||
examples_pool = None
|
||||
else:
|
||||
examples_pool = None
|
||||
|
||||
positives, large_positives = corpus_generator.run(
|
||||
num_samples=num_samples,
|
||||
max_corpus=args.max_corpus,
|
||||
corpus_dir=corpus_dir,
|
||||
doc_length=doc_length,
|
||||
external_path=external_path,
|
||||
source_language=code_language
|
||||
)
|
||||
|
||||
if task_type in ["code_modification_retrieval", "code_comparison_retrieval"]:
|
||||
top1_docs = get_top1([e['text'] for e in positives], args.sim_model_name, [e['text'] for e in large_positives])
|
||||
for i in range(len(top1_docs)):
|
||||
positives[i]['similar'] = top1_docs[i]
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
print("=================== Generate training data ===================")
|
||||
print(f'Task Type: {task_type} | Language: {language} | Code Language: {code_language} | Target Code Language: {tgt_code_language}')
|
||||
start_time = time.time()
|
||||
triplets = gen_triplets(
|
||||
model=model,
|
||||
model_type=model_type,
|
||||
port=port,
|
||||
positives=positives,
|
||||
task_type=task_type,
|
||||
language=language,
|
||||
code_language=code_language,
|
||||
tgt_code_language=tgt_code_language,
|
||||
examples_pool=examples_pool,
|
||||
num_examples=num_examples,
|
||||
thread_count=num_processes,
|
||||
gen_cache_dir=os.path.join(save_dir, language, task_type, "gen_cache_dir"),
|
||||
debug_mode=debug_mode,
|
||||
gen_hard_neg=gen_hard_neg,
|
||||
)
|
||||
save_triplets(
|
||||
triplets=triplets,
|
||||
save_dir=save_dir,
|
||||
task_type=task_type,
|
||||
language=language,
|
||||
code_language=code_language,
|
||||
tgt_code_language=tgt_code_language
|
||||
)
|
||||
end_time = time.time()
|
||||
print("=============================================================")
|
||||
print(f"Time taken: {end_time - start_time:.2f} seconds")
|
||||
print("=============================================================")
|
||||
print("DONE!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = get_args()
|
||||
main(args)
|
||||
@@ -0,0 +1,71 @@
|
||||
from typing import Optional, List
|
||||
|
||||
import faiss
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
from FlagEmbedding import FlagModel
|
||||
|
||||
def create_index(embeddings: np.ndarray, use_gpu: bool = False):
|
||||
index = faiss.IndexFlatIP(len(embeddings[0]))
|
||||
embeddings = np.asarray(embeddings, dtype=np.float32)
|
||||
if use_gpu:
|
||||
co = faiss.GpuMultipleClonerOptions()
|
||||
co.shard = True
|
||||
co.useFloat16 = True
|
||||
index = faiss.index_cpu_to_all_gpus(index, co=co)
|
||||
index.add(embeddings)
|
||||
return index
|
||||
|
||||
|
||||
def search(
|
||||
faiss_index: faiss.Index,
|
||||
k: int = 100,
|
||||
query_embeddings: Optional[np.ndarray] = None,
|
||||
load_path: Optional[str] = None
|
||||
):
|
||||
if query_embeddings is None:
|
||||
query_embeddings = np.load(load_path)
|
||||
|
||||
query_size = len(query_embeddings)
|
||||
|
||||
all_scores = []
|
||||
all_indices = []
|
||||
|
||||
for i in tqdm(range(0, query_size, 32), desc="Searching"):
|
||||
j = min(i + 32, query_size)
|
||||
query_embedding = query_embeddings[i: j]
|
||||
score, indice = faiss_index.search(query_embedding.astype(np.float32), k=k)
|
||||
all_scores.append(score)
|
||||
all_indices.append(indice)
|
||||
|
||||
all_scores = np.concatenate(all_scores, axis=0)
|
||||
all_indices = np.concatenate(all_indices, axis=0)
|
||||
return all_scores, all_indices
|
||||
|
||||
def get_top1(
|
||||
small_docs,
|
||||
encoder_name,
|
||||
docs: List[str],
|
||||
top: int = 1
|
||||
):
|
||||
encoder = FlagModel(encoder_name, trust_remote_code=True)
|
||||
doc_emb = encoder.encode_corpus(docs, max_length=512, batch_size=256)
|
||||
small_doc_emb = encoder.encode_corpus(small_docs, max_length=512, batch_size=256)
|
||||
faiss_index = create_index(doc_emb, True)
|
||||
all_scores, all_indices = search(faiss_index, 1000, small_doc_emb)
|
||||
return_docs = []
|
||||
for i in range(len(all_indices)):
|
||||
return_docs.append([])
|
||||
for idx, score in zip(all_indices[i][20:], all_scores[i][20:]):
|
||||
d1 = set(docs[idx].split())
|
||||
d2 = set(small_docs[i].split())
|
||||
if len(d1 & d2) / len(d1 | d2) > 0.95:
|
||||
continue
|
||||
return_docs[-1].append(docs[idx])
|
||||
if len(return_docs[-1]) >= top:
|
||||
break
|
||||
if len(return_docs[-1]) == 0:
|
||||
print(all_indices[i], all_scores[i])
|
||||
# print(return_docs)
|
||||
del faiss_index
|
||||
return return_docs
|
||||
@@ -0,0 +1,654 @@
|
||||
import os
|
||||
import json
|
||||
import random
|
||||
from tqdm import tqdm
|
||||
from hashlib import md5
|
||||
from warnings import warn
|
||||
from typing import List, Optional
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
from llm import LLM
|
||||
from utils import clean_content
|
||||
from constant import TaskType, Task, SPECIAL_TASK_STEPS, \
|
||||
get_task, get_generation_prompt, get_quality_control_prompt, \
|
||||
get_gen_hard_neg_prompt
|
||||
|
||||
|
||||
def compute_md5(text: str):
|
||||
return md5(text.encode()).hexdigest()
|
||||
|
||||
|
||||
class TripletGenerator(LLM):
|
||||
def __init__(
|
||||
self,
|
||||
model: str = "Qwen2-5-Coder-32B-Instruct",
|
||||
model_type: str = "open-source",
|
||||
port: int = 8000,
|
||||
cache_dir: Optional[str] = None
|
||||
):
|
||||
super().__init__(model, model_type, port)
|
||||
self.cache_dir = cache_dir
|
||||
if self.cache_dir is not None:
|
||||
os.makedirs(self.cache_dir, exist_ok=True)
|
||||
|
||||
def _gen_for_code_modification_retrieval(
|
||||
self,
|
||||
task: Task,
|
||||
text: str,
|
||||
text_b: Optional[str] = None,
|
||||
examples: Optional[List[dict]] = None,
|
||||
debug_mode: bool = False,
|
||||
**kwargs
|
||||
):
|
||||
gen_prompt = get_generation_prompt(
|
||||
task=task,
|
||||
text=text,
|
||||
text_b=text_b,
|
||||
examples=examples,
|
||||
idx=0
|
||||
)
|
||||
response = self.chat(gen_prompt, **kwargs)[0]
|
||||
diff = clean_content(response)
|
||||
gen_prompt = get_generation_prompt(
|
||||
task=task,
|
||||
text=diff,
|
||||
examples=examples,
|
||||
idx=1
|
||||
)
|
||||
response = self.chat(gen_prompt, **kwargs)[0]
|
||||
modification_instr = clean_content(response)
|
||||
|
||||
query = f"{modification_instr}\n```\n{text}\n```"
|
||||
pos = text_b
|
||||
|
||||
if debug_mode:
|
||||
result = {
|
||||
"generation_prompt": gen_prompt,
|
||||
"prompt": task.task_instruction,
|
||||
"query": query,
|
||||
"pos": [pos],
|
||||
"neg": []
|
||||
}
|
||||
else:
|
||||
result = {
|
||||
"prompt": task.task_instruction,
|
||||
"query": query,
|
||||
"pos": [pos],
|
||||
"neg": []
|
||||
}
|
||||
return result
|
||||
|
||||
def _gen_for_code_comparison_retrieval(
|
||||
self,
|
||||
task: Task,
|
||||
text: str,
|
||||
text_b: Optional[str] = None,
|
||||
examples: Optional[List[dict]] = None,
|
||||
debug_mode: bool = False,
|
||||
**kwargs
|
||||
):
|
||||
gen_prompt = get_generation_prompt(
|
||||
task=task,
|
||||
text=text,
|
||||
text_b=text_b,
|
||||
examples=examples,
|
||||
idx=0
|
||||
)
|
||||
response = self.chat(gen_prompt, **kwargs)[0]
|
||||
diff_question = clean_content(response)
|
||||
query = f"{diff_question}\n\nInput Code:\n```\n{text}\n```\n\nOutput Code:\n```\n{text_b}\n```"
|
||||
gen_prompt = get_generation_prompt(
|
||||
task=task,
|
||||
text=query,
|
||||
examples=examples,
|
||||
idx=1
|
||||
)
|
||||
response = self.chat(gen_prompt, **kwargs)[0]
|
||||
pos = clean_content(response)
|
||||
|
||||
if debug_mode:
|
||||
result = {
|
||||
"generation_prompt": gen_prompt,
|
||||
"prompt": task.task_instruction,
|
||||
"query": query,
|
||||
"pos": [pos],
|
||||
"neg": []
|
||||
}
|
||||
else:
|
||||
result = {
|
||||
"prompt": task.task_instruction,
|
||||
"query": query,
|
||||
"pos": [pos],
|
||||
"neg": []
|
||||
}
|
||||
return result
|
||||
|
||||
def _gen_for_code_context_retrieval(
|
||||
self,
|
||||
task: Task,
|
||||
text: str,
|
||||
anchor_points: Optional[tuple] = (0.4, 0.7),
|
||||
**kwargs
|
||||
):
|
||||
former_part, latter_part = self.split_text(
|
||||
text,
|
||||
anchor_points=anchor_points
|
||||
)
|
||||
result = {
|
||||
"prompt": task.task_instruction,
|
||||
"query": former_part,
|
||||
"pos": [latter_part],
|
||||
"neg": []
|
||||
}
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _arrange_query_and_pos(task: Task, input_text: str, response: str):
|
||||
"""
|
||||
Arrange the query and positive example based on the task type.
|
||||
|
||||
Args:
|
||||
- task: Task
|
||||
- input_text: str
|
||||
- response: str
|
||||
|
||||
Returns:
|
||||
- query: str
|
||||
- pos: str
|
||||
"""
|
||||
# TODO: support more task types, including some special task types.
|
||||
if task.main_task_type in ["text2code", "hybrid"]:
|
||||
query = clean_content(response)
|
||||
pos = input_text
|
||||
else:
|
||||
query = input_text
|
||||
pos = clean_content(response)
|
||||
return query, pos
|
||||
|
||||
def _gen_for_normal_task(
|
||||
self,
|
||||
task: Task,
|
||||
text: str,
|
||||
examples: Optional[List[dict]] = None,
|
||||
debug_mode: bool = False,
|
||||
**kwargs
|
||||
):
|
||||
gen_prompt = get_generation_prompt(
|
||||
task=task,
|
||||
text=text,
|
||||
examples=examples
|
||||
)
|
||||
response = self.chat(gen_prompt, **kwargs)[0]
|
||||
|
||||
# Arrange the query and positive example based on the task type.
|
||||
query, pos = self._arrange_query_and_pos(
|
||||
task=task,
|
||||
input_text=text,
|
||||
response=response
|
||||
)
|
||||
|
||||
if debug_mode:
|
||||
result = {
|
||||
"generation_prompt": gen_prompt,
|
||||
"prompt": task.task_instruction,
|
||||
"query": query,
|
||||
"pos": [pos],
|
||||
"neg": [],
|
||||
"response": response
|
||||
}
|
||||
else:
|
||||
result = {
|
||||
"prompt": task.task_instruction,
|
||||
"query": query,
|
||||
"pos": [pos],
|
||||
"neg": []
|
||||
}
|
||||
return result
|
||||
|
||||
def _gen_for_bug_desc_retrieval(
|
||||
self,
|
||||
task: Task,
|
||||
text: str,
|
||||
examples: Optional[List[dict]] = None,
|
||||
debug_mode: bool = False,
|
||||
**kwargs
|
||||
):
|
||||
gen_prompt = get_generation_prompt(
|
||||
task=task,
|
||||
text=text,
|
||||
examples=examples,
|
||||
idx=0
|
||||
)
|
||||
response = self.chat(gen_prompt, **kwargs)[0]
|
||||
if response is None:
|
||||
raise ValueError("Response is None.")
|
||||
buggy_code = response
|
||||
gen_prompt = get_generation_prompt(
|
||||
task=task,
|
||||
text=buggy_code,
|
||||
examples=examples,
|
||||
idx=1
|
||||
)
|
||||
response = self.chat(gen_prompt, **kwargs)[0]
|
||||
query = clean_content(response)
|
||||
pos = text
|
||||
|
||||
if debug_mode:
|
||||
result = {
|
||||
"generation_prompt": gen_prompt,
|
||||
"prompt": task.task_instruction,
|
||||
"query": query,
|
||||
"pos": [pos],
|
||||
"neg": []
|
||||
}
|
||||
else:
|
||||
result = {
|
||||
"prompt": task.task_instruction,
|
||||
"query": query,
|
||||
"pos": [pos],
|
||||
"neg": []
|
||||
}
|
||||
return result
|
||||
|
||||
def _gen_for_two_step_not_use_last(
|
||||
self,
|
||||
task: Task,
|
||||
text: str,
|
||||
examples: Optional[List[dict]] = None,
|
||||
debug_mode: bool = False,
|
||||
reverse_query_pos: bool = False,
|
||||
**kwargs
|
||||
):
|
||||
gen_prompt = get_generation_prompt(
|
||||
task=task,
|
||||
text=text,
|
||||
idx=0
|
||||
)
|
||||
response = self.chat(gen_prompt, **kwargs)[0]
|
||||
query = clean_content(response)
|
||||
gen_prompt = get_generation_prompt(
|
||||
task=task,
|
||||
text=query,
|
||||
examples=examples,
|
||||
idx=1
|
||||
)
|
||||
response = self.chat(gen_prompt, **kwargs)[0]
|
||||
pos = clean_content(response)
|
||||
if reverse_query_pos:
|
||||
query, pos = pos, query
|
||||
|
||||
if debug_mode:
|
||||
result = {
|
||||
"generation_prompt": gen_prompt,
|
||||
"prompt": task.task_instruction,
|
||||
"query": query,
|
||||
"pos": [pos],
|
||||
"neg": []
|
||||
}
|
||||
else:
|
||||
result = {
|
||||
"prompt": task.task_instruction,
|
||||
"query": query,
|
||||
"pos": [pos],
|
||||
"neg": []
|
||||
}
|
||||
return result
|
||||
|
||||
def _gen_for_two_step_use_last(
|
||||
self,
|
||||
task: Task,
|
||||
text: str,
|
||||
examples: Optional[List[dict]] = None,
|
||||
debug_mode: bool = False,
|
||||
reverse_query_pos: bool = False,
|
||||
**kwargs
|
||||
):
|
||||
gen_prompt = get_generation_prompt(
|
||||
task=task,
|
||||
text=text,
|
||||
idx=0
|
||||
)
|
||||
response = self.chat(gen_prompt, **kwargs)[0]
|
||||
query = clean_content(response) + f"\n```\n{text}\n```"
|
||||
gen_prompt = get_generation_prompt(
|
||||
task=task,
|
||||
text=query,
|
||||
examples=examples,
|
||||
idx=1
|
||||
)
|
||||
response = self.chat(gen_prompt, **kwargs)[0]
|
||||
pos = clean_content(response)
|
||||
if reverse_query_pos:
|
||||
query, pos = pos, query
|
||||
|
||||
if debug_mode:
|
||||
result = {
|
||||
"generation_prompt": gen_prompt,
|
||||
"prompt": task.task_instruction,
|
||||
"query": query,
|
||||
"pos": [pos],
|
||||
"neg": []
|
||||
}
|
||||
else:
|
||||
result = {
|
||||
"prompt": task.task_instruction,
|
||||
"query": query,
|
||||
"pos": [pos],
|
||||
"neg": []
|
||||
}
|
||||
return result
|
||||
|
||||
def generate_triplets(
|
||||
self,
|
||||
data: dict,
|
||||
task: Task,
|
||||
examples_pool: Optional[List[dict]] = None,
|
||||
num_examples: int = 3,
|
||||
debug_mode: bool = False,
|
||||
**kwargs
|
||||
):
|
||||
kwargs["remove_thinking"] = not debug_mode
|
||||
|
||||
result_list = []
|
||||
|
||||
examples = None
|
||||
if examples_pool is not None:
|
||||
examples = random.sample(examples_pool, min(num_examples, len(examples_pool)))
|
||||
|
||||
try:
|
||||
if task.task_type in SPECIAL_TASK_STEPS:
|
||||
text = data["text"]
|
||||
|
||||
if task.task_type == TaskType.code_modification_retrieval:
|
||||
text_b = data["similar"][0]
|
||||
|
||||
result = self._gen_for_code_modification_retrieval(
|
||||
task=task,
|
||||
text=text,
|
||||
text_b=text_b,
|
||||
examples=examples,
|
||||
debug_mode=debug_mode
|
||||
)
|
||||
elif task.task_type == TaskType.code_comparison_retrieval:
|
||||
text_b = data["similar"][0]
|
||||
|
||||
result = self._gen_for_code_comparison_retrieval(
|
||||
task=task,
|
||||
text=text,
|
||||
text_b=text_b,
|
||||
examples=examples,
|
||||
debug_mode=debug_mode
|
||||
)
|
||||
elif task.task_type == TaskType.bug_desc_retrieval:
|
||||
result = self._gen_for_bug_desc_retrieval(
|
||||
task=task,
|
||||
text=text,
|
||||
examples=examples,
|
||||
debug_mode=debug_mode
|
||||
)
|
||||
elif task.task_type in [
|
||||
# cf - updated
|
||||
TaskType.code_issue_discussion_retrieval,
|
||||
TaskType.code_version_update_retrieval,
|
||||
TaskType.code_bug_fix_example_retrieval,
|
||||
]:
|
||||
result = self._gen_for_two_step_not_use_last(
|
||||
task=task,
|
||||
text=text,
|
||||
examples=examples,
|
||||
debug_mode=debug_mode,
|
||||
reverse_query_pos=False
|
||||
)
|
||||
elif task.task_type in [
|
||||
# cf - updated
|
||||
TaskType.code_refactoring_pattern_retrieval,
|
||||
TaskType.code_style_guideline_example_retrieval,
|
||||
TaskType.code_migration_retrieval,
|
||||
# jl - updated
|
||||
TaskType.code_optimization_hybrid_retrieval,
|
||||
TaskType.code_best_practices_retrieval,
|
||||
TaskType.security_vulnerability_fix_retrieval,
|
||||
]:
|
||||
result = self._gen_for_two_step_use_last(
|
||||
task=task,
|
||||
text=text,
|
||||
examples=examples,
|
||||
debug_mode=debug_mode,
|
||||
reverse_query_pos=False
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(f"Task type {task.task_type} not implemented.")
|
||||
elif task.task_type == TaskType.code_context_retrieval:
|
||||
text = data["text"]
|
||||
|
||||
result = self._gen_for_code_context_retrieval(
|
||||
task=task,
|
||||
text=text,
|
||||
**kwargs
|
||||
)
|
||||
# NOTE: no need to do quality control for code context retrieval task
|
||||
result_list.append(result)
|
||||
return result_list
|
||||
else:
|
||||
text = data["text"]
|
||||
|
||||
result = self._gen_for_normal_task(
|
||||
task=task,
|
||||
text=text,
|
||||
examples=examples,
|
||||
debug_mode=debug_mode,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
# print(gen_prompt)
|
||||
# print('================================================')
|
||||
qc_prompt = get_quality_control_prompt(
|
||||
task=task,
|
||||
query=result["query"],
|
||||
pos=result["pos"][0]
|
||||
)
|
||||
# print(qc_prompt)
|
||||
# print('*********************************************************************')
|
||||
response = self.chat(qc_prompt, **kwargs)[0]
|
||||
judge = clean_content(response)
|
||||
# print(response, judge)
|
||||
if "1" in judge:
|
||||
if debug_mode:
|
||||
result["judge"] = judge
|
||||
result["judge_response"] = response
|
||||
result_list.append(result)
|
||||
else:
|
||||
if debug_mode:
|
||||
result["judge"] = judge
|
||||
result["judge_response"] = response
|
||||
result_list.append(result)
|
||||
except Exception as e:
|
||||
warn(f"Error: {e}")
|
||||
|
||||
return result_list
|
||||
|
||||
def gen_hard_negatives(self, result: dict, task: Task, num_negatives: int = 7, **kwargs):
|
||||
gen_hard_neg_prompt = get_gen_hard_neg_prompt(
|
||||
task=task,
|
||||
query=result["query"],
|
||||
pos=result["pos"][0]
|
||||
)
|
||||
response_list = self.chat(gen_hard_neg_prompt, n=num_negatives, **kwargs)
|
||||
for response in response_list:
|
||||
if response is None:
|
||||
continue
|
||||
hard_neg = clean_content(response)
|
||||
result["neg"].append(hard_neg)
|
||||
result["neg"] = list(set(result["neg"]))
|
||||
return result
|
||||
|
||||
def run_single(
|
||||
self,
|
||||
data: dict,
|
||||
task: Task,
|
||||
examples_pool: Optional[List[dict]] = None,
|
||||
num_examples: int = 3,
|
||||
debug_mode: bool = False,
|
||||
gen_hard_neg: bool = False,
|
||||
num_negatives: int = 7,
|
||||
**kwargs
|
||||
):
|
||||
result_list = []
|
||||
|
||||
docid = compute_md5(data["text"])
|
||||
if self.cache_dir is not None:
|
||||
gen_data_cache_path = os.path.join(self.cache_dir, f"{docid}.json")
|
||||
if os.path.exists(gen_data_cache_path):
|
||||
with open(gen_data_cache_path, "r", encoding="utf-8") as f:
|
||||
result_list = json.load(f)
|
||||
|
||||
if len(result_list) > 0:
|
||||
if gen_hard_neg:
|
||||
for i in range(len(result_list)):
|
||||
if len(result_list[i]["neg"]) == 0:
|
||||
result_list[i] = self.gen_hard_negatives(
|
||||
result=result_list[i],
|
||||
task=task,
|
||||
num_negatives=num_negatives,
|
||||
**kwargs
|
||||
)
|
||||
# overwrite the cache file
|
||||
with open(gen_data_cache_path, "w", encoding="utf-8") as f:
|
||||
json.dump(result_list, f, indent=4, ensure_ascii=False)
|
||||
return result_list
|
||||
|
||||
triplets = self.generate_triplets(
|
||||
data,
|
||||
task=task,
|
||||
examples_pool=examples_pool,
|
||||
num_examples=num_examples,
|
||||
debug_mode=debug_mode,
|
||||
**kwargs
|
||||
)
|
||||
if len(triplets) == 0:
|
||||
return []
|
||||
|
||||
result = triplets[0]
|
||||
if debug_mode:
|
||||
result["docid"] = docid
|
||||
|
||||
if gen_hard_neg:
|
||||
result = self.gen_hard_negatives(
|
||||
result,
|
||||
task=task,
|
||||
num_negatives=num_negatives,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
result_list.append(result)
|
||||
|
||||
if self.cache_dir is not None:
|
||||
gen_data_cache_path = os.path.join(self.cache_dir, f"{docid}.json")
|
||||
with open(gen_data_cache_path, "w", encoding="utf-8") as f:
|
||||
json.dump(result_list, f, indent=4, ensure_ascii=False)
|
||||
|
||||
return result_list
|
||||
|
||||
def run(
|
||||
self,
|
||||
positives: List[dict],
|
||||
task_type: str,
|
||||
language: str = "en",
|
||||
code_language: str = "python",
|
||||
tgt_code_language: Optional[str] = None,
|
||||
examples_pool: Optional[List[dict]] = None,
|
||||
num_examples: int = 3,
|
||||
tqdm_desc: str = "Generating triplets",
|
||||
debug_mode: bool = False,
|
||||
gen_hard_neg: bool = False,
|
||||
num_negatives: int = 7,
|
||||
thread_count: int = 1,
|
||||
**kwargs
|
||||
):
|
||||
task = get_task(
|
||||
task_type=task_type,
|
||||
language=language,
|
||||
code_language=code_language,
|
||||
tgt_code_language=tgt_code_language
|
||||
)
|
||||
|
||||
result_list = []
|
||||
|
||||
def process_positive(positive):
|
||||
return self.run_single(
|
||||
data=positive,
|
||||
task=task,
|
||||
examples_pool=examples_pool,
|
||||
num_examples=num_examples,
|
||||
debug_mode=debug_mode,
|
||||
gen_hard_neg=gen_hard_neg,
|
||||
num_negatives=num_negatives,
|
||||
**kwargs
|
||||
)
|
||||
# Use thread pool for parallel processing with tqdm progress bar.
|
||||
with ThreadPoolExecutor(max_workers=thread_count) as executor:
|
||||
results = list(tqdm(executor.map(
|
||||
process_positive,
|
||||
positives
|
||||
), total=len(positives), desc=tqdm_desc))
|
||||
|
||||
# Collect results into result_list.
|
||||
for res in results:
|
||||
if isinstance(res, list):
|
||||
result_list.extend(res)
|
||||
else:
|
||||
result_list.append(res)
|
||||
# result_list.extend(results)
|
||||
|
||||
return result_list
|
||||
|
||||
def run_for_gen_neg(
|
||||
self,
|
||||
pairs: List[dict],
|
||||
task_type: str,
|
||||
language: str = "en",
|
||||
code_language: str = "python",
|
||||
tgt_code_language: Optional[str] = None,
|
||||
examples_pool: Optional[List[dict]] = None,
|
||||
num_examples: int = 3,
|
||||
tqdm_desc: str = "Generating triplets",
|
||||
debug_mode: bool = False,
|
||||
gen_hard_neg: bool = False,
|
||||
num_negatives: int = 7,
|
||||
thread_count: int = 1,
|
||||
**kwargs
|
||||
):
|
||||
task = get_task(
|
||||
task_type=task_type,
|
||||
language=language,
|
||||
code_language=code_language,
|
||||
tgt_code_language=tgt_code_language
|
||||
)
|
||||
|
||||
result_list = []
|
||||
|
||||
def gen_single_negative(pair):
|
||||
result = self.gen_hard_negatives(
|
||||
pair,
|
||||
task=task,
|
||||
num_negatives=num_negatives,
|
||||
**kwargs
|
||||
)
|
||||
return [result]
|
||||
|
||||
# Use thread pool for parallel processing with tqdm progress bar.
|
||||
with ThreadPoolExecutor(max_workers=thread_count) as executor:
|
||||
results = list(tqdm(executor.map(
|
||||
gen_single_negative,
|
||||
pairs
|
||||
), total=len(pairs), desc=tqdm_desc))
|
||||
|
||||
# Collect results into result_list.
|
||||
for res in results:
|
||||
if isinstance(res, list):
|
||||
result_list.extend(res)
|
||||
else:
|
||||
result_list.append(res)
|
||||
# result_list.extend(results)
|
||||
|
||||
return result_list
|
||||
@@ -0,0 +1,128 @@
|
||||
import re
|
||||
|
||||
|
||||
def clean_content(content: str):
|
||||
if content is None:
|
||||
raise ValueError("content is None.")
|
||||
|
||||
content = content.split('</think>')[-1].strip('\n').strip()
|
||||
|
||||
if content.startswith('\"') and content.endswith('\"'):
|
||||
content = content[1:-1]
|
||||
|
||||
if content.startswith("```\n") and content.endswith("\n```"):
|
||||
content = content[4:-4]
|
||||
|
||||
return content
|
||||
|
||||
|
||||
def clean_code(code: str, lang: str, length_threshold: int = 30) -> str:
|
||||
cleaned_code = code.strip('\ufeff').strip()
|
||||
if not cleaned_code:
|
||||
return ''
|
||||
|
||||
def clean_empty_lines(text: str) -> str:
|
||||
return re.sub(r'\n\s*\n', '\n', text).strip()
|
||||
|
||||
# 各语言函数/类定义检测正则表达式
|
||||
function_patterns = {
|
||||
"java": r"(?m)^(?!\s*(import|package)\b).*\b(public\s+class|class\s+\w+|void\s+main|new\s+\w+\(|@Override)\b",
|
||||
"python": r"(?m)^(?!\s*(import|from\s+\S+\s+import)\b).*\b(def\s+\w+|class\s+\w+|=\s*\S+|if\s+[:\w]|print\s+)",
|
||||
"javascript": r"(?m)^(?!\s*(import|require\(|export\s)).*\b(function\s+\w+|const\s+\w+|=>|\(\)\s*=>|console\.log)",
|
||||
"php": r"(?m)^(?!\s*(include|require|use)\b).*\b(function\s+\w+|echo\s+\S+|class\s+\w+)",
|
||||
"ruby": r"(?m)^(?!\s*(require|load)\b).*\b(class\s+\w+|def\s+\w+|puts\s+\S+)",
|
||||
"go": r"(?m)^(?!\s*import\b).*\bfunc\s+main\s*\(|type\s+\w+\s+struct",
|
||||
"c#": r"(?m)^(?!\s*using\b).*\b(class\s+\w+|void\s+Main\s*\()",
|
||||
"cplusplus": r"(?m)^(?!#include\b).*\b(int\s+main\s*\(|class\s+\w+|void\s+\w+\s*\(.*\)\s*{)",
|
||||
"c": r"(?m)^(?!#include\b).*\b(int\s+main\s*\(|void\s+\w+\s*\(.*\)\s*{)",
|
||||
"rust": r"(?m)^(?!\s*use\b).*\b(fn\s+main\s*\(|struct\s+\w+|impl\s+\w+)",
|
||||
"typescript": r"(?m)^(?!\s*(import|require\(|export\s)).*\b(interface\s+\w+|class\s+\w+|function\s+\w+)",
|
||||
"perl": r"(?m)^(?!\s*(use|require)\b).*\b(sub\s+\w+|my\s+\$\w+|print\s+\S+)",
|
||||
"shell": r"(?m)^(?!\s*(source|\.)\s).*\b(function\s+\w+|if\s+\[|\$\(|echo\s+\S+)",
|
||||
"sql": r"(?i)\b(CREATE\s+TABLE|SELECT\s+\*|INSERT\s+INTO|UPDATE\s+\w+|DELETE\s+FROM)\b",
|
||||
"batchfile": r"(?m)^(?!\s*@?call\b).*\b(echo\s+\S+|set\s+\w+|if\s+.*\s+==\s+)",
|
||||
"fortran": r"(?mi)^(?!\s*use\b).*\b(program\s+\w+|subroutine\s+\w+|do\s+\d+\s*,\s*\d+)",
|
||||
"haskell": r"(?m)^(?!\s*import\b).*\b(main\s*=\s*do|data\s+\w+|putStrLn\s+\S+)",
|
||||
"lua": r"(?m)^(?!\s*require\b).*\b(function\s+\w+|local\s+\w+|print\s*\()",
|
||||
"powershell": r"(?m)^(?!\s*Import-Module\b).*\b(function\s+\w+|Write-Host\s+\S+|\$\w+\s*=)",
|
||||
"visual_basic": r"(?m)^(?!\s*Imports\b).*\b(Module\s+\w+|Sub\s+Main|Class\s+\w+)"
|
||||
}
|
||||
|
||||
# 各语言注释处理规则
|
||||
comment_patterns = {
|
||||
'java': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
|
||||
'python': (r'#.*?$', re.MULTILINE),
|
||||
'javascript': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
|
||||
'php': (r'//.*?$|#.*?$|/\*.*?\*/|\\/\\/.*?$|#.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
|
||||
'ruby': (r'#.*', re.MULTILINE),
|
||||
'go': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
|
||||
'csharp': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
|
||||
'cplusplus': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
|
||||
'c': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
|
||||
'rust': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
|
||||
'typescript': (r'//.*?$|/\*.*?\*/|\\/\\/.*?$|\\/\*.*?\*\\/', re.DOTALL | re.MULTILINE),
|
||||
'perl': (r'#.*', re.MULTILINE),
|
||||
'shell': (r'#.*', re.MULTILINE),
|
||||
'sql': (r'--.*?$|/\*.*?\*/', re.DOTALL),
|
||||
'batchfile': (r'^\s*(REM|@REM|::).*', re.MULTILINE | re.IGNORECASE),
|
||||
'fortran': (r'!.*', re.MULTILINE),
|
||||
'haskell': (r'--.*', re.MULTILINE),
|
||||
'lua': (r'--.*?$|--\[\[.*?\]\]', re.DOTALL),
|
||||
'powershell': (r'<#.*?#>|#.*', re.DOTALL),
|
||||
'visual_basic': (r"'.*", re.MULTILINE),
|
||||
}
|
||||
|
||||
# 执行注释清理
|
||||
if lang in comment_patterns:
|
||||
pattern, flags = comment_patterns[lang]
|
||||
cleaned_code = re.sub(pattern, '', cleaned_code, flags=flags)
|
||||
cleaned_code = clean_empty_lines(cleaned_code)
|
||||
|
||||
# 特殊语言处理规则
|
||||
if lang == 'fortran':
|
||||
cleaned_code = re.sub(r'^[Cc*].*', '', cleaned_code, flags=re.MULTILINE)
|
||||
elif lang == 'sql':
|
||||
cleaned_code = re.sub(r'/\*.*?\*/', '', cleaned_code, flags=re.DOTALL)
|
||||
elif lang == 'python':
|
||||
cleaned_code = re.sub(r'^\s*#.*', '', cleaned_code, flags=re.MULTILINE)
|
||||
|
||||
# 函数定义检测及内容验证
|
||||
def has_valid_code(text: str, lang: str) -> bool:
|
||||
pattern = function_patterns.get(lang)
|
||||
if not pattern:
|
||||
return len(text.strip()) > 0
|
||||
|
||||
# 增强检测逻辑
|
||||
if lang == 'batchfile':
|
||||
return bool(re.search(r'^\s*@?echo\b|:\w+', text, re.MULTILINE))
|
||||
if lang == 'shell':
|
||||
return bool(re.search(r'^\s*(if|for|while|case|echo|export|shopt|source)\b', text, re.MULTILINE))
|
||||
if lang == 'python':
|
||||
if re.search(r'^\s*(def|class)\s+\w+', text, re.MULTILINE):
|
||||
return bool(re.search(r'^\s+[^\s#]', text, re.MULTILINE))
|
||||
return False
|
||||
if lang == 'ruby':
|
||||
return bool(re.search(r'(def\s+\w+|class\s+\w+).*?\n\s+[^\s#]', text, re.MULTILINE))
|
||||
return bool(re.search(pattern, text, re.DOTALL | re.MULTILINE))
|
||||
|
||||
# 最终有效性检查
|
||||
if not has_valid_code(cleaned_code, lang):
|
||||
return ''
|
||||
|
||||
cleaned_code = cleaned_code.strip('\ufeff').strip()
|
||||
|
||||
if len(cleaned_code) < length_threshold:
|
||||
return ''
|
||||
|
||||
return cleaned_code
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_text = "\/\/ ----------------------------------------------------------------------\n\/\/ ----------------------------------------------------------------------\n\/\/\n\/\/ File: StrMaxProjection.h\n\/\/ Author: mgrosso \n\/\/ Created: Mon Jul 17 14:39:22 PDT 2006 on caliban\n\/\/ Project: \n\/\/ Purpose: \n\/\/ \n\/\/ $Id$\n\/\/ ----------------------------------------------------------------------\n\/\/ ----------------------------------------------------------------------\n\n#ifndef STRMAXPROJECTION_H\n#define STRMAXPROJECTION_H 1\n\n#include \"StrMinProjection.h\"\n\nclass StrMaxProjection : public StrMinProjection\n{\n public:\n StrMaxProjection(ExpressionPtr &operand);\n virtual ~StrMaxProjection();\n virtual AbstractProjectionPtr copy();\n\n protected:\n int compare(const char *lhs, const char *rhs);\n\n private:\n \/\/not implemented\n StrMaxProjection();\n StrMaxProjection( const StrMaxProjection &rhs );\n StrMaxProjection &operator=( const StrMaxProjection &rhs );\n};\n\n#endif \/* STRMAXPROJECTION_H *\/"
|
||||
|
||||
result = clean_code(test_text, "c", 200)
|
||||
print(result)
|
||||
|
||||
test_text = "\/**\n * Copyright (c) Microsoft Corporation. All rights reserved.\n * Licensed under the MIT License. See License.txt in the project root for\n * license information.\n *\n * Code generated by Microsoft (R) AutoRest Code Generator.\n *\/\n\npackage com.microsoft.azure.management.datafactory.v2018_06_01;\n\nimport com.fasterxml.jackson.annotation.JsonProperty;\nimport com.fasterxml.jackson.annotation.JsonTypeInfo;\nimport com.fasterxml.jackson.annotation.JsonTypeName;\n\n\/**\n * The location of Google Cloud Storage dataset.\n *\/\n@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.PROPERTY, property = \"type\", defaultImpl = GoogleCloudStorageLocation.class)\n@JsonTypeName(\"GoogleCloudStorageLocation\")\npublic class GoogleCloudStorageLocation extends DatasetLocation {\n \/**\n * Specify the bucketName of Google Cloud Storage. Type: string (or\n * Expression with resultType string).\n *\/\n @JsonProperty(value = \"bucketName\")\n private Object bucketName;\n\n \/**\n * Specify the version of Google Cloud Storage. Type: string (or Expression\n * with resultType string).\n *\/\n @JsonProperty(value = \"version\")\n private Object version;\n\n \/**\n * Get specify the bucketName of Google Cloud Storage. Type: string (or Expression with resultType string).\n *\n * @return the bucketName value\n *\/\n public Object bucketName() {\n return this.bucketName;\n }\n\n \/**\n * Set specify the bucketName of Google Cloud Storage. Type: string (or Expression with resultType string).\n *\n * @param bucketName the bucketName value to set\n * @return the GoogleCloudStorageLocation object itself.\n *\/\n public GoogleCloudStorageLocation withBucketName(Object bucketName) {\n this.bucketName = bucketName;\n return this;\n }\n\n \/**\n * Get specify the version of Google Cloud Storage. Type: string (or Expression with resultType string).\n *\n * @return the version value\n *\/\n public Object version() {\n return this.version;\n }\n\n \/**\n * Set specify the version of Google Cloud Storage. Type: string (or Expression with resultType string).\n *\n * @param version the version value to set\n * @return the GoogleCloudStorageLocation object itself.\n *\/\n public GoogleCloudStorageLocation withVersion(Object version) {\n this.version = version;\n return this;\n }\n\n}"
|
||||
result = clean_code(test_text, "java", 200)
|
||||
print(result)
|
||||
@@ -0,0 +1,22 @@
|
||||
cd ./code-rag-bench/retrieval/
|
||||
|
||||
output_dir='result'
|
||||
|
||||
for dataset_name in "humaneval" "mbpp" "repoeval" "ds1000_all_completion" "odex_en" "swe-bench-lite"
|
||||
do
|
||||
echo "dataset_name: ${dataset_name}"
|
||||
python main.py \
|
||||
--embedder_name_or_path BAAI/bge-code-v1 \
|
||||
--embedder_model_class decoder-only-base \
|
||||
--query_instruction_format_for_retrieval '<instruct>{}\n<query>{}' \
|
||||
--embedder_query_max_length 2048 \
|
||||
--embedder_passage_max_length 2048 \
|
||||
--trust_remote_code True \
|
||||
--pooling_method last_token \
|
||||
--embedder_batch_size 64 \
|
||||
--devices cuda:0 cuda:1 cuda:2 cuda:3 cuda:4 cuda:5 cuda:6 cuda:7 \
|
||||
--cache_dir ./cache \
|
||||
--dataset $dataset_name \
|
||||
--output_file ../../${output_dir}/${dataset_name}_output.json \
|
||||
--results_file ../../${output_dir}/${dataset_name}_results.json
|
||||
done
|
||||
@@ -0,0 +1,7 @@
|
||||
cd ./code-rag-bench/retrieval/
|
||||
|
||||
for dataset_name in "humaneval" "mbpp" "live_code_bench" "ds1000" "odex" "repoeval_repo" "swebench_repo"
|
||||
do
|
||||
echo "dataset_name: ${dataset_name}"
|
||||
PYTHONPATH=./ python create/${dataset_name}.py
|
||||
done
|
||||
@@ -0,0 +1,34 @@
|
||||
from typing import List
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from FlagEmbedding.abc.evaluation import (
|
||||
AbsEvalModelArgs as CodeRAGEvalModelArgs,
|
||||
)
|
||||
|
||||
@dataclass
|
||||
class CodeRAGEvalArgs:
|
||||
dataset: str = field(
|
||||
default='humaneval',
|
||||
metadata={
|
||||
"help": "Task to evaluate. Default: humaneval. Available tasks: "
|
||||
"['humaneval', 'mbpp', 'live_code_bench', 'ds1000', 'odex', 'repoeval_repo', 'swebench_repo', 'code_search_net']",
|
||||
}
|
||||
)
|
||||
max_length: int = field(
|
||||
default=2048, metadata={"help": "Max length to use for evaluation."}
|
||||
)
|
||||
batch_size: int = field(
|
||||
default=64, metadata={"help": "Batch size for evaluation."}
|
||||
)
|
||||
output_file: str = field(
|
||||
default="outputs.json",
|
||||
metadata={
|
||||
"help": "Specify the filepath if you want to save the retrieval (evaluation) results."
|
||||
},
|
||||
)
|
||||
results_file: str = field(
|
||||
default="results.json",
|
||||
metadata={
|
||||
"help": "Specify the filepath if you want to save the retrieval results."
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,52 @@
|
||||
import argparse
|
||||
import datasets
|
||||
import os
|
||||
from tqdm import tqdm
|
||||
import random
|
||||
from create.utils import save_tsv_dict, save_file_jsonl, load_jsonlines
|
||||
|
||||
def document2code(data, split="train"):
|
||||
data = data[split]
|
||||
code_search_net_data_queries = []
|
||||
code_search_net_data_docs = []
|
||||
code_search_net_data_qrels = []
|
||||
|
||||
for item in tqdm(data):
|
||||
doc = item["func_documentation_string"]
|
||||
code = item["func_code_string"]
|
||||
doc_id = "{repository_name}_{func_path_in_repository}_{func_name}_doc".format_map(item)
|
||||
code_id = "{repository_name}_{func_path_in_repository}_{func_name}_code".format_map(item)
|
||||
code_search_net_data_queries.append({"_id": doc_id, "text": doc, "metadata": {}})
|
||||
code_search_net_data_docs.append({"_id": code_id, "title": item["func_name"], "text": code, "metadata": {}})
|
||||
code_search_net_data_qrels.append({"query-id": doc_id, "corpus-id": code_id, "score": 1})
|
||||
|
||||
return code_search_net_data_queries, code_search_net_data_docs, code_search_net_data_qrels
|
||||
|
||||
def main():
|
||||
#### /print debug information to stdout
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--language", type=str, default="python", help="codesearch net language")
|
||||
parser.add_argument("--output_dir", type=str, default="datasets")
|
||||
|
||||
args = parser.parse_args()
|
||||
dataset = datasets.load_dataset("code_search_net", args.language)
|
||||
|
||||
path = os.path.join(args.output_dir, "code_search_net_{}".format(args.language))
|
||||
os.makedirs(path)
|
||||
os.makedirs(os.path.join(path, "qrels"))
|
||||
|
||||
docs = []
|
||||
queries = []
|
||||
for split in ["train", "validation", "test"]:
|
||||
queries_split, docs_split, qrels_split = document2code(dataset, split)
|
||||
docs += docs_split
|
||||
queries += queries_split
|
||||
|
||||
save_tsv_dict(qrels_split, os.path.join(path, "qrels", "{}.tsv".format(split)), ["query-id", "corpus-id", "score"])
|
||||
|
||||
save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
|
||||
save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,123 @@
|
||||
import io
|
||||
import os
|
||||
import fcntl
|
||||
import pathlib
|
||||
import zipfile
|
||||
import argparse
|
||||
import requests
|
||||
import warnings
|
||||
import itertools
|
||||
from tqdm import tqdm
|
||||
from datasets import load_dataset
|
||||
from create.utils import save_tsv_dict, save_file_jsonl
|
||||
|
||||
|
||||
# Load dataset
|
||||
def download_source(source_dir):
|
||||
src = source_dir / "ds1000.py"
|
||||
url = "https://github.com/HKUNLP/DS-1000/blob/49c1c543ada8b58138181333cdc62e613204efcf/ds1000.py?raw=true"
|
||||
lock = src.with_suffix(".lock")
|
||||
with open(lock, "w") as f_lock:
|
||||
fcntl.flock(f_lock, fcntl.LOCK_EX)
|
||||
if not src.exists():
|
||||
warnings.warn(f"DS-1000 source is being saved to {src}.")
|
||||
print("Downloading source code...")
|
||||
r = requests.get(url, stream=True)
|
||||
with open(src, "wb") as f_src:
|
||||
f_src.write(r.content)
|
||||
open(src.parent / "__init__.py", "w").close()
|
||||
print("Done.")
|
||||
fcntl.flock(f_lock, fcntl.LOCK_UN)
|
||||
|
||||
def download_dataset(source_dir):
|
||||
path = source_dir / "ds1000_data"
|
||||
url = "https://github.com/HKUNLP/DS-1000/blob/49c1c543ada8b58138181333cdc62e613204efcf/ds1000_data.zip?raw=true"
|
||||
lock = path.with_suffix(".lock")
|
||||
with open(lock, "w") as f_lock:
|
||||
fcntl.flock(f_lock, fcntl.LOCK_EX)
|
||||
if not path.exists():
|
||||
warnings.warn(f"DS-1000 data is being saved to {path}.")
|
||||
print("Downloading dataset...")
|
||||
r = requests.get(url, stream=True)
|
||||
z = zipfile.ZipFile(io.BytesIO(r.content))
|
||||
z.extractall(source_dir)
|
||||
print("Done.")
|
||||
fcntl.flock(f_lock, fcntl.LOCK_UN)
|
||||
|
||||
def get_dataset(source_dir, mode: str = "Completion", key: str = "All"):
|
||||
"""Returns dataset for the task or an iterable of any object, that get_prompt can handle"""
|
||||
from ds.ds1000 import DS1000Dataset
|
||||
|
||||
data = DS1000Dataset(source_dir / "ds1000_data", mode=mode).data
|
||||
if key == "All":
|
||||
if mode == "Insertion":
|
||||
warnings.warn(
|
||||
"Insertion not supported for Matplotlib. Only running others."
|
||||
)
|
||||
data = {k: v for k, v in data.items() if k != "Matplotlib"}
|
||||
dataset = list(itertools.chain(*data.values()))
|
||||
else:
|
||||
dataset = data[key]
|
||||
return dataset
|
||||
|
||||
|
||||
# Collect queries, docs, and relations
|
||||
def document2code(data: list):
|
||||
queries, docs, qrels = [], [], []
|
||||
|
||||
# collect doc corpus
|
||||
code_docs = load_dataset("neulab/docprompting-conala", "docs")["train"]
|
||||
for i in range(len(code_docs)):
|
||||
docs.append({
|
||||
"_id": str(i),
|
||||
"title": code_docs[i]["doc_id"],
|
||||
"text": code_docs[i]["doc_content"],
|
||||
"metadata": {}
|
||||
})
|
||||
|
||||
# load canonical docs
|
||||
ds1000 = load_dataset("json", data_files={"test": args.canonical_file})["test"]
|
||||
for idx,item in enumerate(tqdm(data)):
|
||||
example = item.data
|
||||
query = example["prompt"]
|
||||
query_id = f"{example['lib']}_{example['perturbation_origin_id']}"
|
||||
queries.append({"_id": query_id, "text": query, "metadata": {}})
|
||||
|
||||
doc_ids = [doc["title"] for doc in ds1000[idx]["docs"]]
|
||||
for doc_id in doc_ids:
|
||||
corpus_id = code_docs["doc_id"].index(doc_id)
|
||||
corpus_id = str(corpus_id)
|
||||
qrels.append({"query-id": query_id, "corpus-id": corpus_id, "score": 1})
|
||||
|
||||
return queries, docs, qrels
|
||||
|
||||
|
||||
def main():
|
||||
args.source_dir = pathlib.Path(__file__).parent.parent / args.source_dir
|
||||
os.makedirs(args.source_dir, exist_ok=True)
|
||||
download_source(args.source_dir)
|
||||
download_dataset(args.source_dir)
|
||||
dataset = get_dataset(args.source_dir, mode=args.mode, key=args.key)
|
||||
|
||||
path = os.path.join(args.output_dir, f"ds1000_{args.key.lower()}_{args.mode.lower()}")
|
||||
os.makedirs(path, exist_ok=True)
|
||||
os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
|
||||
|
||||
queries, docs, qrels = document2code(dataset)
|
||||
save_tsv_dict(qrels, os.path.join(path, "qrels", "test.tsv"), ["query-id", "corpus-id", "score"])
|
||||
|
||||
save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
|
||||
save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--source_dir", type=str, default="ds")
|
||||
parser.add_argument("--output_dir", type=str, default="datasets")
|
||||
parser.add_argument("--mode", type=str, default="Completion", choices=["Completion", "Insertion"])
|
||||
parser.add_argument("--key", type=str, default="All",
|
||||
choices=["All", "Numpy", "Pandas", "Scipy", "Matplotlib", "Sklearn", "Tensorflow", "Pytorch"])
|
||||
parser.add_argument("--canonical_file", type=str, default="datasets/canonical/ds1000_docs.json")
|
||||
args = parser.parse_args()
|
||||
|
||||
main()
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Aggregate all code-generation datasets."""
|
||||
|
||||
import os
|
||||
import json
|
||||
import datasets
|
||||
import argparse
|
||||
from create.utils import save_tsv_dict
|
||||
from create.humaneval import document2code as d2c_humaneval
|
||||
from create.mbpp import document2code as d2c_mbpp
|
||||
|
||||
D2C_FUNC_DICT = {
|
||||
"humaneval": d2c_humaneval,
|
||||
"mbpp": d2c_mbpp,
|
||||
}
|
||||
SPLIT_DICT = {
|
||||
"humaneval": ["test"],
|
||||
"mbpp": ["train", "test", "validation", "prompt"],
|
||||
}
|
||||
HF_NAME_DICT = {
|
||||
"humaneval": "openai_humaneval",
|
||||
"mbpp": "mbpp",
|
||||
}
|
||||
|
||||
|
||||
def save_file_jsonl(data, path):
|
||||
with open(path,'w') as fw:
|
||||
for item in data:
|
||||
fw.write(json.dumps(item) + '\n')
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--dataset_names", type=str, nargs='+', default=["humaneval", "mbpp"])
|
||||
parser.add_argument("--output_dir", type=str, default="datasets")
|
||||
parser.add_argument("--output_name", type=str, default="general-programming")
|
||||
args = parser.parse_args()
|
||||
|
||||
path = os.path.join(args.output_dir, args.output_name)
|
||||
os.makedirs(path)
|
||||
os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
|
||||
|
||||
split_dict = {}
|
||||
for dataset_name in args.dataset_names:
|
||||
for split in SPLIT_DICT[dataset_name]:
|
||||
if split not in split_dict:
|
||||
split_dict[split] = []
|
||||
split_dict[split].append(dataset_name)
|
||||
|
||||
dataset_dict = {
|
||||
dataset_name: datasets.load_dataset(HF_NAME_DICT[dataset_name])
|
||||
for dataset_name in args.dataset_names
|
||||
}
|
||||
docs, queries = [], []
|
||||
for split, ds_names in split_dict.items():
|
||||
for ds in ds_names:
|
||||
dataset = dataset_dict[ds]
|
||||
|
||||
queries_split, docs_split, qrels_split = D2C_FUNC_DICT[ds](dataset, split)
|
||||
docs += docs_split
|
||||
queries += queries_split
|
||||
|
||||
qrels_path = os.path.join(path, "qrels", f"{split}.tsv")
|
||||
save_tsv_dict(qrels_split, qrels_path, ["query-id", "corpus-id", "score"])
|
||||
|
||||
save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
|
||||
save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,58 @@
|
||||
import os
|
||||
import argparse
|
||||
import datasets
|
||||
from tqdm import tqdm
|
||||
from create.utils import save_tsv_dict, save_file_jsonl
|
||||
|
||||
|
||||
def document2code(data, split="test"):
|
||||
data = data[split]
|
||||
queries, docs, qrels = [], [], []
|
||||
|
||||
for item in tqdm(data):
|
||||
doc = item["prompt"]
|
||||
code = item["prompt"] + '\n' + item["canonical_solution"]
|
||||
doc_id = "{task_id}_doc".format_map(item)
|
||||
code_id = "{task_id}_code".format_map(item)
|
||||
|
||||
queries.append({"_id": doc_id, "text": doc, "metadata": {}})
|
||||
docs.append({"_id": code_id, "title": item["entry_point"], "text": code, "metadata": {}})
|
||||
qrels.append({"query-id": doc_id, "corpus-id": code_id, "score": 1})
|
||||
|
||||
return queries, docs, qrels
|
||||
|
||||
|
||||
def main():
|
||||
dataset = datasets.load_dataset(args.dataset_name)
|
||||
|
||||
path = os.path.join(args.output_dir, args.output_name)
|
||||
os.makedirs(path, exist_ok=True)
|
||||
os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
|
||||
|
||||
queries, docs, qrels = document2code(dataset, split="test")
|
||||
save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
|
||||
save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
|
||||
qrels_path = os.path.join(path, "qrels", "test.tsv")
|
||||
save_tsv_dict(qrels, qrels_path, ["query-id", "corpus-id", "score"])
|
||||
|
||||
# create canonical file if not existent yet
|
||||
if not os.path.exists(args.canonical_file):
|
||||
canonical_solutions = []
|
||||
for doc in docs:
|
||||
canonical_solutions.append([{
|
||||
"text": doc["text"], "title": doc["title"]
|
||||
}])
|
||||
canonical_dataset = dataset["test"].add_column("docs", canonical_solutions)
|
||||
canonical_dataset.to_json(args.canonical_file)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--dataset_name", type=str, default="openai_humaneval")
|
||||
parser.add_argument("--output_name", type=str, default="humaneval")
|
||||
parser.add_argument("--canonical_file", type=str,
|
||||
default="datasets/canonical/humaneval_solutions.json")
|
||||
parser.add_argument("--output_dir", type=str, default="datasets")
|
||||
args = parser.parse_args()
|
||||
|
||||
main()
|
||||
@@ -0,0 +1,53 @@
|
||||
import os
|
||||
import argparse
|
||||
import datasets
|
||||
from tqdm import tqdm
|
||||
from datasets import load_dataset
|
||||
from create.utils import save_tsv_dict, save_file_jsonl
|
||||
|
||||
|
||||
def get_queries(data, split="test") -> list[dict]:
|
||||
queries = [{
|
||||
"_id": item["question_id"] + '__' + item["contest_id"],
|
||||
"text": item["question_content"],
|
||||
"metadata": {}
|
||||
} for item in data[split]]
|
||||
return queries
|
||||
|
||||
def get_corpus(hf_name: str, cache_dir: str) -> list[dict]:
|
||||
dataset = load_dataset(hf_name, cache_dir=cache_dir)["train"]
|
||||
corpus = [
|
||||
{"_id": i, "text": item["text"], "title": item["title"]}
|
||||
for i,item in enumerate(dataset)
|
||||
]
|
||||
return corpus
|
||||
|
||||
|
||||
def main():
|
||||
dataset = datasets.load_dataset(args.dataset_name, cache_dir=args.cache_dir)
|
||||
|
||||
path = os.path.join(args.output_dir, args.output_name)
|
||||
os.makedirs(path, exist_ok=True)
|
||||
os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
|
||||
|
||||
queries = get_queries(dataset, split="test")
|
||||
save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
|
||||
|
||||
docs = get_corpus(args.corpus_name, args.cache_dir)
|
||||
save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
|
||||
|
||||
qrels = [] # no ground-truth solutions
|
||||
qrels_path = os.path.join(path, "qrels", "test.tsv")
|
||||
save_tsv_dict(qrels, qrels_path, ["query-id", "corpus-id", "score"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--dataset_name", type=str, default="livecodebench/code_generation")
|
||||
parser.add_argument("--corpus_name", type=str, default="code-rag-bench/programming-solutions")
|
||||
parser.add_argument("--cache_dir", type=str, default="/scratch/zhiruow/data")
|
||||
parser.add_argument("--output_name", type=str, default="livecodebench")
|
||||
parser.add_argument("--output_dir", type=str, default="datasets")
|
||||
args = parser.parse_args()
|
||||
|
||||
main()
|
||||
@@ -0,0 +1,78 @@
|
||||
import os
|
||||
import argparse
|
||||
import datasets
|
||||
from tqdm import tqdm
|
||||
from create.utils import save_tsv_dict, save_file_jsonl
|
||||
|
||||
|
||||
def get_function_name(code: str) -> str:
|
||||
"""Parse the function name for a code snippet string."""
|
||||
lines = code.split('\n')
|
||||
for line in lines:
|
||||
if line.lstrip().startswith("def "):
|
||||
break
|
||||
func_name = line.lstrip()[4: ]
|
||||
func_name = func_name.split('(')[0]
|
||||
return func_name
|
||||
|
||||
|
||||
def document2code(data, split="test"):
|
||||
data = data[split]
|
||||
queries, docs, qrels = [], [], []
|
||||
|
||||
for item in tqdm(data):
|
||||
doc = item["text"]
|
||||
code = "# " + item["text"] + '\n' + item["code"]
|
||||
doc_id = "{task_id}_doc".format_map(item)
|
||||
code_id = "{task_id}_code".format_map(item)
|
||||
|
||||
queries.append({"_id": doc_id, "text": doc, "metadata": {}})
|
||||
docs.append({"_id": code_id, "title": get_function_name(item["code"]), "text": code, "metadata": {}})
|
||||
qrels.append({"query-id": doc_id, "corpus-id": code_id, "score": 1})
|
||||
|
||||
return queries, docs, qrels
|
||||
|
||||
|
||||
def main():
|
||||
dataset = datasets.load_dataset(args.dataset_name)
|
||||
|
||||
path = os.path.join(args.output_dir, args.output_name)
|
||||
os.makedirs(path, exist_ok=True)
|
||||
os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
|
||||
|
||||
docs, queries = [], []
|
||||
for split in args.splits:
|
||||
queries_split, docs_split, qrels_split = document2code(dataset, split)
|
||||
docs += docs_split
|
||||
queries += queries_split
|
||||
|
||||
qrels_path = os.path.join(path, "qrels", f"{split}.tsv")
|
||||
save_tsv_dict(qrels_split, qrels_path, ["query-id", "corpus-id", "score"])
|
||||
|
||||
# create canonical file for test split if not existent yet
|
||||
if split == "test" and (not os.path.exists(args.canonical_file)):
|
||||
canonical_solutions = []
|
||||
for doc in docs_split:
|
||||
canonical_solutions.append([{
|
||||
"text": doc["text"], "title": doc["title"]
|
||||
}])
|
||||
canonical_dataset = dataset["test"].add_column("docs", canonical_solutions)
|
||||
canonical_dataset.to_json(args.canonical_file)
|
||||
|
||||
save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
|
||||
save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--dataset_name", type=str, default="google-research-datasets/mbpp")
|
||||
# parser.add_argument("--dataset_name", type=str, default="code-rag-bench/mbpp")
|
||||
parser.add_argument("--splits", type=str, default=["train", "validation", "test"],
|
||||
choices=["train", "validation", "test", "prompt"])
|
||||
parser.add_argument("--output_name", type=str, default="mbpp")
|
||||
parser.add_argument("--output_dir", type=str, default="datasets")
|
||||
parser.add_argument("--canonical_file", type=str,
|
||||
default="datasets/canonical/mbpp_solutions.json")
|
||||
|
||||
args = parser.parse_args()
|
||||
main()
|
||||
@@ -0,0 +1,76 @@
|
||||
import os
|
||||
import re
|
||||
import random
|
||||
import argparse
|
||||
import datasets
|
||||
from tqdm import tqdm
|
||||
from collections import Counter
|
||||
from datasets import load_dataset
|
||||
from create.utils import save_tsv_dict, save_file_jsonl
|
||||
|
||||
|
||||
def document2code(data, split="test"):
|
||||
data = data[split]
|
||||
queries, docs, qrels = [], [], []
|
||||
|
||||
# build doc corpus
|
||||
code_docs = load_dataset("neulab/docprompting-conala", "docs")["train"]
|
||||
for i in range(len(code_docs)):
|
||||
docs.append({
|
||||
"_id": str(i),
|
||||
"title": code_docs[i]["doc_id"],
|
||||
"text": code_docs[i]["doc_content"],
|
||||
"metadata": {}
|
||||
})
|
||||
|
||||
# load canonical docs
|
||||
odex = load_dataset("json", data_files={"test": args.canonical_file})["test"]
|
||||
# collect queries and query-doc matching
|
||||
for idx,item in enumerate(tqdm(data)):
|
||||
query = item["intent"]
|
||||
query_id = f"{idx}_{item['task_id']}"
|
||||
queries.append({"_id": query_id, "text": query, "metadata": {}})
|
||||
|
||||
doc_ids = [doc["title"] for doc in odex[idx]["docs"]]
|
||||
for doc_id in doc_ids:
|
||||
corpus_id = code_docs["doc_id"].index(doc_id)
|
||||
corpus_id = str(corpus_id)
|
||||
qrels.append({"query-id": query_id, "corpus-id": corpus_id, "score": 1})
|
||||
|
||||
return queries, docs, qrels
|
||||
|
||||
|
||||
def main():
|
||||
if '_' in args.dataset_name:
|
||||
dataset_name = args.dataset_name.split('_')[0]
|
||||
language = args.dataset_name.split('_')[1]
|
||||
else:
|
||||
dataset_name = args.dataset_name
|
||||
language = 'en'
|
||||
dataset = datasets.load_dataset(dataset_name, language) # english version by default
|
||||
|
||||
path = os.path.join(args.output_dir, args.output_name.replace('en', language))
|
||||
os.makedirs(path, exist_ok=True)
|
||||
os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
|
||||
|
||||
docs, queries = [], []
|
||||
for split in ["test"]:
|
||||
queries_split, docs_split, qrels_split = document2code(dataset, split)
|
||||
docs += docs_split
|
||||
queries += queries_split
|
||||
|
||||
save_tsv_dict(qrels_split, os.path.join(path, "qrels", "{}.tsv".format(split)), ["query-id", "corpus-id", "score"])
|
||||
|
||||
save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
|
||||
save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--dataset_name", type=str, default="neulab/odex")
|
||||
parser.add_argument("--output_name", type=str, default="odex_en")
|
||||
parser.add_argument("--canonical_file", type=str, default="datasets/canonical/odex_docs.json")
|
||||
parser.add_argument("--output_dir", type=str, default="datasets")
|
||||
args = parser.parse_args()
|
||||
|
||||
main()
|
||||
@@ -0,0 +1,306 @@
|
||||
import io
|
||||
import os
|
||||
import glob
|
||||
import json
|
||||
import argparse
|
||||
import requests
|
||||
import zipfile
|
||||
from collections import defaultdict
|
||||
from create.utils import save_tsv_dict, save_file_jsonl
|
||||
|
||||
REPOs_line_and_api = [
|
||||
'huggingface_diffusers',
|
||||
'nerfstudio-project_nerfstudio',
|
||||
'awslabs_fortuna',
|
||||
'huggingface_evaluate',
|
||||
'google_vizier',
|
||||
'alibaba_FederatedScope',
|
||||
'pytorch_rl',
|
||||
'opendilab_ACE',
|
||||
]
|
||||
|
||||
REPOs_function = [
|
||||
"amazon-science_patchcore-inspection",
|
||||
"deepmind_tracr",
|
||||
"facebookresearch_omnivore",
|
||||
"google_lightweight_mmm",
|
||||
"lucidrains_imagen-pytorch",
|
||||
"maxhumber_redframes",
|
||||
]
|
||||
|
||||
REPO_DIRs = {
|
||||
"api": "repositories/line_and_api_level",
|
||||
"line": "repositories/line_and_api_level",
|
||||
"function": "repositories/function_level",
|
||||
}
|
||||
|
||||
|
||||
def iterate_repository(base_dir: str, repo: str) -> dict:
|
||||
pattern = os.path.join(f'{base_dir}/{repo}', "**", "*.py")
|
||||
files = glob.glob(pattern, recursive=True)
|
||||
|
||||
skipped_files = []
|
||||
loaded_code_files = dict()
|
||||
base_dir_list = os.path.normpath(base_dir).split(os.sep)
|
||||
for fname in files:
|
||||
try:
|
||||
code = open(fname, 'r', encoding='utf8').read()
|
||||
fpath_tuple = tuple(os.path.normpath(fname).split(os.sep)[len(base_dir_list):])
|
||||
loaded_code_files[fpath_tuple]= code
|
||||
except Exception as e:
|
||||
skipped_files.append((fname, e))
|
||||
continue
|
||||
|
||||
if len(skipped_files) > 0:
|
||||
print(f"Skipped {len(skipped_files)} out of {len(files)} files due to I/O errors")
|
||||
for fname, e in skipped_files:
|
||||
print(f"{fname}: {e}")
|
||||
return loaded_code_files
|
||||
|
||||
|
||||
def window_overlap(span: tuple, target_span: tuple) -> bool:
|
||||
if span[0] >= target_span[1] or span[1] <= target_span[0]:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class RepoWindowMaker:
|
||||
def __init__(self, base_dir, repo, tasks, window_size, slice_size):
|
||||
self.base_dir = base_dir
|
||||
self.repo = repo
|
||||
self.window_size = window_size
|
||||
self.slice_size = slice_size
|
||||
self.slice_step = 1 if window_size // slice_size == 0 else window_size // slice_size
|
||||
self.tasks = tasks
|
||||
self.source_code_files = iterate_repository(base_dir, repo)
|
||||
|
||||
def _buid_windows_for_a_file(self, fpath_tuple, code):
|
||||
code_windows = []
|
||||
code_lines = code.splitlines()
|
||||
delta_size = self.window_size // 2
|
||||
for line_no in range(0, len(code_lines), self.slice_step): # line_no starts from 0
|
||||
start_line_no = max(0, line_no - delta_size)
|
||||
end_line_no = min(len(code_lines), line_no + self.window_size - delta_size)
|
||||
window_lines = [i for i in code_lines[start_line_no:end_line_no]]
|
||||
if not window_lines: # all empty lines
|
||||
continue
|
||||
window_text = '\n'.join(window_lines)
|
||||
code_windows.append({
|
||||
'context': window_text,
|
||||
'metadata': {
|
||||
'fpath_tuple': fpath_tuple,
|
||||
'line_no': line_no,
|
||||
'start_line_no': start_line_no,
|
||||
'end_line_no': end_line_no,
|
||||
'window_size': self.window_size,
|
||||
'repo': self.repo,
|
||||
'slice_size': self.slice_size,
|
||||
}
|
||||
})
|
||||
return code_windows
|
||||
|
||||
def _merge_windows_with_same_context(self, code_windows):
|
||||
merged_code_windows = defaultdict(list)
|
||||
for code_window in code_windows:
|
||||
context = code_window['context']
|
||||
metadata = code_window['metadata']
|
||||
merged_code_windows[context].append(metadata)
|
||||
json_lines = []
|
||||
for context, metadata_list in merged_code_windows.items():
|
||||
json_lines.append({
|
||||
'context': context,
|
||||
'metadata': metadata_list
|
||||
})
|
||||
return json_lines
|
||||
|
||||
def build_windows(self):
|
||||
all_code_windows = []
|
||||
for fpath_tuple, code in self.source_code_files.items():
|
||||
all_code_windows += self._buid_windows_for_a_file(fpath_tuple, code)
|
||||
merged_code_windows = self._merge_windows_with_same_context(all_code_windows)
|
||||
print(f'build {len(merged_code_windows)} windows for {self.repo} with window size {self.window_size} and slice {self.slice_size}')
|
||||
ground_truth_indices = {}
|
||||
for task in self.tasks:
|
||||
fpath_tuple = tuple(task['metadata']['fpath_tuple'])
|
||||
line_no = task['metadata']['line_no']
|
||||
start_line_no = task['metadata']['context_start_lineno']
|
||||
for i, window in enumerate(merged_code_windows):
|
||||
if window["metadata"][0]["fpath_tuple"] != fpath_tuple:
|
||||
continue
|
||||
if any([
|
||||
window_overlap(
|
||||
(sub_window["start_line_no"], sub_window["end_line_no"]),
|
||||
(start_line_no, line_no + 1)
|
||||
)
|
||||
for sub_window in window["metadata"]
|
||||
]):
|
||||
if i not in ground_truth_indices:
|
||||
ground_truth_indices[i] = []
|
||||
ground_truth_indices[i].append(task["metadata"]["task_id"])
|
||||
|
||||
return merged_code_windows, ground_truth_indices
|
||||
|
||||
|
||||
def download_data(directory: str = "repoeval"):
|
||||
os.makedirs(directory, exist_ok=True)
|
||||
|
||||
datasets_dir = os.path.join(directory, "datasets")
|
||||
repos_lineapi_dir = os.path.join(directory, "repositories", "line_and_api_level")
|
||||
repos_function_dir = os.path.join(directory, "repositories", "function_level")
|
||||
|
||||
print(f"Start downloading the necessary `datasets` and `repositories` files.")
|
||||
if not os.path.exists(datasets_dir):
|
||||
print(f"Start downloading the `datasets`.")
|
||||
datasets_url = "https://github.com/microsoft/CodeT/raw/main/RepoCoder/datasets/datasets.zip"
|
||||
r = requests.get(datasets_url, stream=True)
|
||||
z = zipfile.ZipFile(io.BytesIO(r.content))
|
||||
z.extractall(datasets_dir)
|
||||
print("Finished downloading the `datasets` files.")
|
||||
|
||||
if not os.path.exists(repos_lineapi_dir):
|
||||
print(f"Start downloading the `repositories` (line_and_api).")
|
||||
repos_lineapi_url = "https://github.com/microsoft/CodeT/raw/main/RepoCoder/repositories/line_and_api_level.zip"
|
||||
r = requests.get(repos_lineapi_url, stream=True)
|
||||
z = zipfile.ZipFile(io.BytesIO(r.content))
|
||||
z.extractall(repos_lineapi_dir)
|
||||
|
||||
if not os.path.exists(repos_function_dir):
|
||||
print(f"Start downloading the `repositories` (function).")
|
||||
# repos_function_url = "https://github.com/microsoft/CodeT/raw/main/RepoCoder/repositories/function_level.zip"
|
||||
repos_function_url = "https://github.com/Veronicium/repoeval_debug/raw/main/function_level.zip"
|
||||
r = requests.get(repos_function_url, stream=True)
|
||||
z = zipfile.ZipFile(io.BytesIO(r.content))
|
||||
z.extractall(repos_function_dir)
|
||||
print("Finished downloading the `repositories` files.")
|
||||
|
||||
|
||||
def repo2code(
|
||||
repo: str, data_cache_dir: str,
|
||||
split: str, context_length: str,
|
||||
window_size: int, slice_size: int
|
||||
):
|
||||
# load test examples
|
||||
file_name = f"{split}_level_completion_{context_length}_context_codex.test.jsonl"
|
||||
if split == 'function':
|
||||
file_name = file_name.replace('.test.jsonl', '.test.clean.jsonl')
|
||||
|
||||
task_path = os.path.join(data_cache_dir, "datasets", file_name)
|
||||
tasks = [json.loads(l.rstrip()) for l in open(task_path, 'r')]
|
||||
tasks = [task for task in tasks if repo == task['metadata']['task_id'].split('/')[0]]
|
||||
|
||||
# collect queries
|
||||
queries = []
|
||||
for task in tasks:
|
||||
query_id = task["metadata"]["task_id"]
|
||||
# text = '\n'.join(task["prompt"].split('\n')[-2:])
|
||||
text = task["prompt"]
|
||||
metadata = task["metadata"]
|
||||
queries.append({"_id": query_id, "text": text, "metadata": metadata})
|
||||
|
||||
base_dir = os.path.join(data_cache_dir, REPO_DIRs[split])
|
||||
repo_window_maker = RepoWindowMaker(base_dir, repo, tasks, window_size, slice_size)
|
||||
windows, ground_truth_indices = repo_window_maker.build_windows()
|
||||
corpus, qrels = [], []
|
||||
query_id2gt = {task['metadata']['task_id']:[] for task in tasks}
|
||||
for i, window in enumerate(windows):
|
||||
path = '-'.join(window["metadata"][0]["fpath_tuple"])
|
||||
line = f"{window['metadata'][0]['start_line_no']}-{window['metadata'][-1]['end_line_no']}"
|
||||
corpus_id = f"{repo}_{path}_{line}"
|
||||
corpus.append({
|
||||
"_id": corpus_id, "title": path,
|
||||
"text": window["context"], "metadata": window["metadata"]
|
||||
})
|
||||
if i in ground_truth_indices:
|
||||
for query_id in ground_truth_indices[i]:
|
||||
qrels.append({"query-id": query_id, "corpus-id": corpus_id, "score": 1})
|
||||
query_id2gt[query_id].append({"title": corpus_id.replace('_', '/'), "text": window["context"]})
|
||||
|
||||
return queries, corpus, qrels, query_id2gt
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--output_dir", type=str, default="datasets")
|
||||
parser.add_argument("--results_dir", type=str, default="results")
|
||||
parser.add_argument("--split", type=str, required=True, choices=["api", "line", "function"])
|
||||
parser.add_argument("--context_length", type=str, default="1k", choices=["1k", "2k", "4k"])
|
||||
parser.add_argument("--data_cache_dir", type=str, default="output/repoeval")
|
||||
parser.add_argument("--window_size", type=int, default=20)
|
||||
parser.add_argument("--slice_size", type=int, default=2)
|
||||
args = parser.parse_args()
|
||||
|
||||
download_data(args.data_cache_dir)
|
||||
|
||||
path = os.path.join(args.output_dir, "repoeval", args.split)
|
||||
os.makedirs(path, exist_ok=True)
|
||||
REPOs = REPOs_function if args.split == "function" else REPOs_line_and_api
|
||||
|
||||
file_name = f"{args.split}_level_completion_{args.context_length}_context_codex.test.jsonl"
|
||||
data_path = os.path.join(args.data_cache_dir, "datasets", file_name)
|
||||
data = [json.loads(l.rstrip()) for l in open(data_path, 'r')]
|
||||
|
||||
# preprocess function completion data (the data in the RepoCoder repo isn't correctly formatted)
|
||||
if args.split == 'function':
|
||||
repo2idx = {}
|
||||
for task in data:
|
||||
repo = task['metadata']['task_id'].replace('--', '_').split('/')[0]
|
||||
if repo not in repo2idx:
|
||||
repo2idx[repo] = 0
|
||||
task['metadata']['task_id'] = task['metadata']['task_id'].replace('--', '_').replace('idx', str(repo2idx[repo]))
|
||||
task['metadata']['line_no'] = task['metadata']['lineno']
|
||||
repo2idx[repo] += 1
|
||||
|
||||
new_data_path = data_path.replace('.test.jsonl', '.test.clean.jsonl')
|
||||
with open(new_data_path, 'w') as f:
|
||||
for task in data:
|
||||
repo = task['metadata']['task_id'].split('/')[0]
|
||||
if repo not in REPOs:
|
||||
continue
|
||||
f.write(json.dumps(task) + '\n')
|
||||
|
||||
data = [json.loads(l.rstrip()) for l in open(new_data_path, 'r')]
|
||||
|
||||
# build query, docs, and qrels for each repository
|
||||
queries, corpus, qrels = [], [], []
|
||||
query_id2gt = {}
|
||||
for repo in REPOs:
|
||||
repo_queries, repo_corpus, repo_qrels, repo_query_id2gt = repo2code(
|
||||
repo, args.data_cache_dir,
|
||||
args.split, args.context_length,
|
||||
args.window_size, args.slice_size
|
||||
)
|
||||
queries += repo_queries
|
||||
corpus += repo_corpus
|
||||
qrels += repo_qrels
|
||||
query_id2gt.update(repo_query_id2gt)
|
||||
|
||||
save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
|
||||
save_file_jsonl(corpus, os.path.join(path, "corpus.jsonl"))
|
||||
save_tsv_dict(qrels, os.path.join(path, "qrels", "test.tsv"), ["query-id", "corpus-id", "score"])
|
||||
|
||||
gt_data = []
|
||||
for example in data:
|
||||
query_id = example['metadata']['task_id']
|
||||
gt = query_id2gt[query_id]
|
||||
new_example = {
|
||||
"prompt": example["prompt"],
|
||||
"reference": example["metadata"]["ground_truth"],
|
||||
"docs": gt[:10],
|
||||
"metadata": {k:v for k,v in example["metadata"].items() if k != "ground_truth"},
|
||||
}
|
||||
gt_data.append(new_example)
|
||||
|
||||
results_file = os.path.join(args.results_dir, f"repoeval-{args.split}-{args.context_length}-gt.jsonl")
|
||||
with open(results_file, "w") as fw:
|
||||
for ex in gt_data:
|
||||
fw.write(json.dumps(ex) + "\n")
|
||||
|
||||
results_file = os.path.join(args.results_dir, f"repoeval-{args.split}-{args.context_length}-infile.jsonl")
|
||||
with open(results_file, "w") as fw:
|
||||
for ex in gt_data:
|
||||
ex = {k:v for k,v in ex.items() if k != "docs"}
|
||||
ex["docs"] = []
|
||||
fw.write(json.dumps(ex) + "\n")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,306 @@
|
||||
import io
|
||||
import os
|
||||
import glob
|
||||
import json
|
||||
import argparse
|
||||
import requests
|
||||
import zipfile
|
||||
from collections import defaultdict
|
||||
from create.utils import save_tsv_dict, save_file_jsonl
|
||||
|
||||
REPOs_line_and_api = [
|
||||
'huggingface_diffusers',
|
||||
'nerfstudio-project_nerfstudio',
|
||||
'awslabs_fortuna',
|
||||
'huggingface_evaluate',
|
||||
'google_vizier',
|
||||
'alibaba_FederatedScope',
|
||||
'pytorch_rl',
|
||||
'opendilab_ACE',
|
||||
]
|
||||
|
||||
REPOs_function = [
|
||||
"amazon-science_patchcore-inspection",
|
||||
"deepmind_tracr",
|
||||
"facebookresearch_omnivore",
|
||||
"google_lightweight_mmm",
|
||||
"lucidrains_imagen-pytorch",
|
||||
"maxhumber_redframes",
|
||||
]
|
||||
|
||||
REPO_DIRs = {
|
||||
"api": "repositories/line_and_api_level",
|
||||
"line": "repositories/line_and_api_level",
|
||||
"function": "repositories/function_level",
|
||||
}
|
||||
|
||||
|
||||
def iterate_repository(base_dir: str, repo: str) -> dict:
|
||||
pattern = os.path.join(f'{base_dir}/{repo}', "**", "*.py")
|
||||
files = glob.glob(pattern, recursive=True)
|
||||
|
||||
skipped_files = []
|
||||
loaded_code_files = dict()
|
||||
base_dir_list = os.path.normpath(base_dir).split(os.sep)
|
||||
for fname in files:
|
||||
try:
|
||||
code = open(fname, 'r', encoding='utf8').read()
|
||||
fpath_tuple = tuple(os.path.normpath(fname).split(os.sep)[len(base_dir_list):])
|
||||
loaded_code_files[fpath_tuple]= code
|
||||
except Exception as e:
|
||||
skipped_files.append((fname, e))
|
||||
continue
|
||||
|
||||
if len(skipped_files) > 0:
|
||||
print(f"Skipped {len(skipped_files)} out of {len(files)} files due to I/O errors")
|
||||
for fname, e in skipped_files:
|
||||
print(f"{fname}: {e}")
|
||||
return loaded_code_files
|
||||
|
||||
|
||||
def window_overlap(span: tuple, target_span: tuple) -> bool:
|
||||
if span[0] >= target_span[1] or span[1] <= target_span[0]:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class RepoWindowMaker:
|
||||
def __init__(self, base_dir, repo, tasks, window_size, slice_size):
|
||||
self.base_dir = base_dir
|
||||
self.repo = repo
|
||||
self.window_size = window_size
|
||||
self.slice_size = slice_size
|
||||
self.slice_step = 1 if window_size // slice_size == 0 else window_size // slice_size
|
||||
self.tasks = tasks
|
||||
self.source_code_files = iterate_repository(base_dir, repo)
|
||||
|
||||
def _buid_windows_for_a_file(self, fpath_tuple, code):
|
||||
code_windows = []
|
||||
code_lines = code.splitlines()
|
||||
delta_size = self.window_size // 2
|
||||
for line_no in range(0, len(code_lines), self.slice_step): # line_no starts from 0
|
||||
start_line_no = max(0, line_no - delta_size)
|
||||
end_line_no = min(len(code_lines), line_no + self.window_size - delta_size)
|
||||
window_lines = [i for i in code_lines[start_line_no:end_line_no]]
|
||||
if not window_lines: # all empty lines
|
||||
continue
|
||||
window_text = '\n'.join(window_lines)
|
||||
code_windows.append({
|
||||
'context': window_text,
|
||||
'metadata': {
|
||||
'fpath_tuple': fpath_tuple,
|
||||
'line_no': line_no,
|
||||
'start_line_no': start_line_no,
|
||||
'end_line_no': end_line_no,
|
||||
'window_size': self.window_size,
|
||||
'repo': self.repo,
|
||||
'slice_size': self.slice_size,
|
||||
}
|
||||
})
|
||||
return code_windows
|
||||
|
||||
def _merge_windows_with_same_context(self, code_windows):
|
||||
merged_code_windows = defaultdict(list)
|
||||
for code_window in code_windows:
|
||||
context = code_window['context']
|
||||
metadata = code_window['metadata']
|
||||
merged_code_windows[context].append(metadata)
|
||||
json_lines = []
|
||||
for context, metadata_list in merged_code_windows.items():
|
||||
json_lines.append({
|
||||
'context': context,
|
||||
'metadata': metadata_list
|
||||
})
|
||||
return json_lines
|
||||
|
||||
def build_windows(self):
|
||||
all_code_windows = []
|
||||
for fpath_tuple, code in self.source_code_files.items():
|
||||
all_code_windows += self._buid_windows_for_a_file(fpath_tuple, code)
|
||||
merged_code_windows = self._merge_windows_with_same_context(all_code_windows)
|
||||
print(f'build {len(merged_code_windows)} windows for {self.repo} with window size {self.window_size} and slice {self.slice_size}')
|
||||
ground_truth_indices = {}
|
||||
for task in self.tasks:
|
||||
fpath_tuple = tuple(task['metadata']['fpath_tuple'])
|
||||
line_no = task['metadata']['line_no']
|
||||
start_line_no = task['metadata']['context_start_lineno']
|
||||
for i, window in enumerate(merged_code_windows):
|
||||
# print(window["metadata"][0]["fpath_tuple"], fpath_tuple)
|
||||
if window["metadata"][0]["fpath_tuple"] != fpath_tuple and ' '.join(list(window["metadata"][0]["fpath_tuple"])) != ' '.join(list(fpath_tuple)):
|
||||
continue
|
||||
# print(1)
|
||||
if any([
|
||||
window_overlap(
|
||||
(sub_window["start_line_no"], sub_window["end_line_no"]),
|
||||
(start_line_no, line_no + 1)
|
||||
)
|
||||
for sub_window in window["metadata"]
|
||||
]):
|
||||
# print('test')
|
||||
if i not in ground_truth_indices:
|
||||
ground_truth_indices[i] = []
|
||||
ground_truth_indices[i].append(task["metadata"]["task_id"])
|
||||
# sys.exit()
|
||||
return merged_code_windows, ground_truth_indices
|
||||
|
||||
|
||||
def download_data(directory: str = "repoeval"):
|
||||
os.makedirs(directory, exist_ok=True)
|
||||
|
||||
datasets_dir = os.path.join(directory, "datasets")
|
||||
repos_function_dir = os.path.join(directory, "repositories", "function_level")
|
||||
|
||||
print(f"Start downloading the necessary `datasets` and `repositories` files.")
|
||||
if not os.path.exists(datasets_dir):
|
||||
print(f"Start downloading the `datasets`.")
|
||||
datasets_url = "https://github.com/microsoft/CodeT/raw/main/RepoCoder/datasets/datasets.zip"
|
||||
r = requests.get(datasets_url, stream=True)
|
||||
z = zipfile.ZipFile(io.BytesIO(r.content))
|
||||
z.extractall(datasets_dir)
|
||||
print("Finished downloading the `datasets` files.")
|
||||
|
||||
import shutil
|
||||
shutil.rmtree(repos_function_dir)
|
||||
if not os.path.exists(repos_function_dir):
|
||||
print(f"Start downloading the `repositories` (function).")
|
||||
repos_function_url = "https://github.com/microsoft/CodeT/raw/main/RepoCoder/repositories/function_level.zip"
|
||||
# repos_function_url = "https://github.com/Veronicium/repoeval_debug/raw/main/function_level.zip"
|
||||
r = requests.get(repos_function_url, stream=True)
|
||||
z = zipfile.ZipFile(io.BytesIO(r.content))
|
||||
z.extractall(repos_function_dir)
|
||||
print("Finished downloading the `repositories` files.")
|
||||
|
||||
|
||||
def repo2code(
|
||||
repo: str, tasks: list[dict], data_cache_dir: str,
|
||||
split: str, context_length: str,
|
||||
window_size: int, slice_size: int
|
||||
):
|
||||
# collect queries
|
||||
queries = []
|
||||
for task in tasks:
|
||||
query_id = task["metadata"]["task_id"]
|
||||
# text = '\n'.join(task["prompt"].split('\n')[-2:])
|
||||
text = task["prompt"]
|
||||
metadata = task["metadata"]
|
||||
queries.append({"_id": query_id, "text": text, "metadata": metadata})
|
||||
|
||||
base_dir = os.path.join(data_cache_dir, REPO_DIRs[split])
|
||||
repo_window_maker = RepoWindowMaker(base_dir, repo, tasks, window_size, slice_size)
|
||||
windows, ground_truth_indices = repo_window_maker.build_windows()
|
||||
corpus, qrels = [], []
|
||||
query_id2gt = {task['metadata']['task_id']:[] for task in tasks}
|
||||
for i, window in enumerate(windows):
|
||||
path = '-'.join(window["metadata"][0]["fpath_tuple"])
|
||||
line = f"{window['metadata'][0]['start_line_no']}-{window['metadata'][-1]['end_line_no']}"
|
||||
corpus_id = f"{repo}_{path}_{line}"
|
||||
corpus.append({
|
||||
"_id": corpus_id, "title": path,
|
||||
"text": window["context"], "metadata": window["metadata"]
|
||||
})
|
||||
# print(windows, ground_truth_indices)
|
||||
if i in ground_truth_indices:
|
||||
for query_id in ground_truth_indices[i]:
|
||||
qrels.append({"query-id": query_id, "corpus-id": corpus_id, "score": 1})
|
||||
query_id2gt[query_id].append({"title": corpus_id.replace('_', '/'), "text": window["context"]})
|
||||
|
||||
return queries, corpus, qrels, query_id2gt
|
||||
|
||||
|
||||
def main():
|
||||
download_data(args.data_cache_dir)
|
||||
|
||||
REPOs = REPOs_function if args.split == "function" else REPOs_line_and_api
|
||||
|
||||
file_name = f"{args.split}_level_completion_{args.context_length}_context_codex.test.jsonl"
|
||||
data_path = os.path.join(args.data_cache_dir, "datasets", file_name)
|
||||
data = [json.loads(l.rstrip()) for l in open(data_path, 'r')]
|
||||
|
||||
# preprocess function completion data (the data in the RepoCoder repo isn't correctly formatted)
|
||||
if args.split == 'function':
|
||||
repo2idx = {}
|
||||
for task in data:
|
||||
repo = task['metadata']['task_id'].replace('--', '_').split('/')[0]
|
||||
if repo not in repo2idx:
|
||||
repo2idx[repo] = 0
|
||||
task['metadata']['task_id'] = task['metadata']['task_id'].replace('--', '_').replace('idx', str(repo2idx[repo]))
|
||||
task['metadata']['line_no'] = task['metadata']['lineno']
|
||||
repo2idx[repo] += 1
|
||||
|
||||
new_data_path = data_path.replace('.test.jsonl', '.test.clean.jsonl')
|
||||
with open(new_data_path, 'w') as f:
|
||||
for task in data:
|
||||
repo = task['metadata']['task_id'].split('/')[0]
|
||||
if repo not in REPOs:
|
||||
continue
|
||||
f.write(json.dumps(task) + '\n')
|
||||
|
||||
data = [json.loads(l.rstrip()) for l in open(new_data_path, 'r')]
|
||||
|
||||
# group data instances by repository
|
||||
data_dict = {}
|
||||
for ex in data:
|
||||
repo_name = ex["metadata"]["task_id"]
|
||||
repo_name = repo_name.split('/')[0]
|
||||
if repo_name not in data_dict:
|
||||
data_dict[repo_name] = []
|
||||
data_dict[repo_name].append(ex)
|
||||
|
||||
# build query, docs, and qrels for each repository
|
||||
for repo in REPOs:
|
||||
queries, corpus, qrels, query_id2gt = repo2code(
|
||||
repo, data_dict[repo], args.data_cache_dir,
|
||||
args.split, args.context_length,
|
||||
args.window_size, args.slice_size
|
||||
)
|
||||
print(len(queries))
|
||||
if len(qrels) == 0:
|
||||
print(repo)
|
||||
# sys.exit()
|
||||
continue
|
||||
# sys.exit()
|
||||
path = os.path.join(args.output_dir, f"repoeval__{repo}")
|
||||
os.makedirs(path, exist_ok=True)
|
||||
save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
|
||||
save_file_jsonl(corpus, os.path.join(path, "corpus.jsonl"))
|
||||
save_tsv_dict(qrels, os.path.join(path, "qrels", "test.tsv"), ["query-id", "corpus-id", "score"])
|
||||
|
||||
gt_data = []
|
||||
for example in data_dict[repo]:
|
||||
query_id = example['metadata']['task_id']
|
||||
gt = query_id2gt[query_id]
|
||||
new_example = {
|
||||
"prompt": example["prompt"],
|
||||
"reference": example["metadata"]["ground_truth"],
|
||||
"docs": gt[:10],
|
||||
"metadata": {k:v for k,v in example["metadata"].items() if k != "ground_truth"},
|
||||
}
|
||||
gt_data.append(new_example)
|
||||
|
||||
os.makedirs(args.results_dir, exist_ok=True)
|
||||
|
||||
results_file = os.path.join(args.results_dir, f"repoeval-{args.split}-{repo}-{args.context_length}-gt.jsonl")
|
||||
with open(results_file, "w") as fw:
|
||||
for ex in gt_data:
|
||||
fw.write(json.dumps(ex) + "\n")
|
||||
|
||||
results_file = os.path.join(args.results_dir, f"repoeval-{args.split}-{repo}-{args.context_length}-infile.jsonl")
|
||||
with open(results_file, "w") as fw:
|
||||
for ex in gt_data:
|
||||
ex = {k:v for k,v in ex.items() if k != "docs"}
|
||||
ex["docs"] = []
|
||||
fw.write(json.dumps(ex) + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--output_dir", type=str, default="datasets")
|
||||
parser.add_argument("--results_dir", type=str, default="results")
|
||||
parser.add_argument("--split", type=str, default="function", choices=["function"])
|
||||
parser.add_argument("--context_length", type=str, default="2k", choices=["1k", "2k", "4k"])
|
||||
parser.add_argument("--data_cache_dir", type=str, default="output/repoeval")
|
||||
parser.add_argument("--window_size", type=int, default=50)
|
||||
parser.add_argument("--slice_size", type=int, default=5)
|
||||
args = parser.parse_args()
|
||||
|
||||
main()
|
||||
@@ -0,0 +1,247 @@
|
||||
import os
|
||||
import re
|
||||
import chardet
|
||||
import unidiff
|
||||
import argparse
|
||||
import datasets
|
||||
import traceback
|
||||
import subprocess
|
||||
from git import Repo
|
||||
from tqdm import tqdm
|
||||
from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
from create.utils import save_tsv_dict, save_file_jsonl
|
||||
|
||||
# %% Get oracle file contents
|
||||
|
||||
# get oracle file contents from the repo
|
||||
class ContextManager:
|
||||
def __init__(self, repo_path, base_commit, verbose=False):
|
||||
self.repo_path = Path(repo_path).resolve().as_posix()
|
||||
self.old_dir = os.getcwd()
|
||||
self.base_commit = base_commit
|
||||
self.verbose = verbose
|
||||
|
||||
def __enter__(self):
|
||||
os.chdir(self.repo_path)
|
||||
cmd = f"git reset --hard {self.base_commit} && git clean -fdxq"
|
||||
if self.verbose:
|
||||
subprocess.run(cmd, shell=True, check=True)
|
||||
else:
|
||||
subprocess.run(
|
||||
cmd,
|
||||
shell=True,
|
||||
check=True,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
return self
|
||||
|
||||
def get_environment(self):
|
||||
raise NotImplementedError() # TODO: activate conda environment and return the environment file
|
||||
|
||||
def get_readme_files(self):
|
||||
files = os.listdir(self.repo_path)
|
||||
files = list(filter(lambda x: os.path.isfile(x), files))
|
||||
files = list(filter(lambda x: x.lower().startswith("readme"), files))
|
||||
return files
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
os.chdir(self.old_dir)
|
||||
|
||||
|
||||
class AutoContextManager(ContextManager):
|
||||
"""Automatically clones the repo if it doesn't exist"""
|
||||
|
||||
def __init__(self, instance, root_dir=None, verbose=False, token=None):
|
||||
if token is None:
|
||||
token = os.environ.get("GITHUB_TOKEN", "git")
|
||||
self.tempdir = None
|
||||
if root_dir is None:
|
||||
self.tempdir = TemporaryDirectory()
|
||||
root_dir = self.tempdir.name
|
||||
self.root_dir = root_dir
|
||||
repo_dir = os.path.join(self.root_dir, instance["repo"].replace("/", "__"))
|
||||
if not os.path.exists(repo_dir):
|
||||
repo_url = (
|
||||
f"https://{token}@github.com/swe-bench/"
|
||||
+ instance["repo"].replace("/", "__")
|
||||
+ ".git"
|
||||
)
|
||||
if verbose:
|
||||
print(f"Cloning {instance['repo']} to {root_dir}")
|
||||
Repo.clone_from(repo_url, repo_dir)
|
||||
super().__init__(repo_dir, instance["base_commit"], verbose=verbose)
|
||||
self.instance = instance
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
if self.tempdir is not None:
|
||||
self.tempdir.cleanup()
|
||||
return super().__exit__(exc_type, exc_val, exc_tb)
|
||||
|
||||
|
||||
def ingest_files(filenames):
|
||||
files_dict = dict()
|
||||
for filename in filenames:
|
||||
with open(filename) as f:
|
||||
content = f.read()
|
||||
files_dict[filename] = content
|
||||
return files_dict
|
||||
|
||||
def get_oracle_filenames(instance):
|
||||
"""
|
||||
Returns the filenames that are changed in the patch
|
||||
"""
|
||||
source_files = {
|
||||
patch_file.source_file.split("a/", 1)[-1]
|
||||
for patch_file in unidiff.PatchSet(instance["patch"])
|
||||
}
|
||||
gold_docs = set()
|
||||
for source_file in source_files:
|
||||
gold_docs.add(source_file)
|
||||
return gold_docs
|
||||
|
||||
|
||||
# get all file contents from the repo
|
||||
def is_test(name, test_phrases=None):
|
||||
if test_phrases is None:
|
||||
test_phrases = ["test", "tests", "testing"]
|
||||
words = set(re.split(r" |_|\/|\.", name.lower()))
|
||||
return any(word in words for word in test_phrases)
|
||||
|
||||
def list_files(root_dir, include_tests=False):
|
||||
files = []
|
||||
for filename in Path(root_dir).rglob("*.py"):
|
||||
if not include_tests and is_test(filename.as_posix()):
|
||||
continue
|
||||
files.append(filename.relative_to(root_dir).as_posix())
|
||||
return files
|
||||
|
||||
def detect_encoding(filename):
|
||||
"""
|
||||
Detect the encoding of a file
|
||||
"""
|
||||
with open(filename, "rb") as file:
|
||||
rawdata = file.read()
|
||||
return chardet.detect(rawdata)["encoding"]
|
||||
|
||||
def ingest_directory_contents(root_dir, include_tests=False):
|
||||
files_content = {}
|
||||
for relative_path in list_files(root_dir, include_tests=include_tests):
|
||||
filename = os.path.join(root_dir, relative_path)
|
||||
encoding = detect_encoding(filename)
|
||||
if encoding is None:
|
||||
content = "[BINARY DATA FILE]"
|
||||
else:
|
||||
try:
|
||||
with open(filename, encoding=encoding) as file:
|
||||
content = file.read()
|
||||
except (UnicodeDecodeError, LookupError):
|
||||
content = "[BINARY DATA FILE]"
|
||||
files_content[relative_path] = content
|
||||
return files_content
|
||||
|
||||
def get_file_contents(input_instances, verbose: bool = False, tmp_dir: str = "/scratch"):
|
||||
orig_dir = os.getcwd()
|
||||
with TemporaryDirectory(dir=tmp_dir if os.path.exists(tmp_dir) else "/tmp") as root_dir:
|
||||
for instance_id, instance in tqdm(
|
||||
input_instances.items(),
|
||||
total=len(input_instances),
|
||||
desc="Getting file contents",
|
||||
):
|
||||
try:
|
||||
with AutoContextManager(instance, root_dir, verbose=verbose) as cm:
|
||||
readmes = cm.get_readme_files()
|
||||
instance["readmes"] = ingest_files(readmes)
|
||||
instance["oracle_file_contents"] = ingest_files(get_oracle_filenames(instance))
|
||||
instance["file_contents"] = ingest_directory_contents(cm.repo_path)
|
||||
assert all([
|
||||
okey in instance["file_contents"]
|
||||
for okey in instance["oracle_file_contents"].keys()
|
||||
])
|
||||
except Exception as e:
|
||||
print(f"Failed on instance {instance_id}", e)
|
||||
traceback.print_exc()
|
||||
finally:
|
||||
# if AutoContextManager fails to exit properly future exits will return the wrong directory
|
||||
os.chdir(orig_dir)
|
||||
os.chdir(orig_dir)
|
||||
|
||||
|
||||
# %% Get queries, docs, and qrels
|
||||
|
||||
def document2code(data, split: str = "test"):
|
||||
subset = data[split]
|
||||
if args.num_examples is not None:
|
||||
import random
|
||||
indices = random.sample([i for i in range(len(subset))], args.num_examples)
|
||||
subset = subset.select(indices)
|
||||
print(subset)
|
||||
|
||||
# get queries for each example
|
||||
queries = [
|
||||
{
|
||||
"_id": item["instance_id"],
|
||||
"text": item["problem_statement"],
|
||||
"metadata": {}
|
||||
}
|
||||
for item in subset
|
||||
]
|
||||
|
||||
subset_dict = {x["instance_id"]: x for x in subset}
|
||||
get_file_contents(subset_dict, tmp_dir=args.tmp_dir)
|
||||
|
||||
# collect all docs, i.e., code chunks from the repo
|
||||
docs = []
|
||||
for instance_id, instance in subset_dict.items():
|
||||
print(f"Instance #{instance_id}: {len(instance['oracle_file_contents'])} oracle / {len(instance['file_contents'])} files")
|
||||
for filename, content in instance["file_contents"].items():
|
||||
docs.append({
|
||||
"_id": f"{instance_id}_{filename}",
|
||||
"title": filename,
|
||||
"text": content,
|
||||
"metadata": {},
|
||||
})
|
||||
|
||||
# find ground-truth docs for each example
|
||||
qrels = []
|
||||
for instance_id, instance in subset_dict.items():
|
||||
for filename, content in instance["oracle_file_contents"].items():
|
||||
qrels.append({
|
||||
"query-id": instance_id,
|
||||
"corpus-id": f"{instance_id}_{filename}",
|
||||
"score": 1
|
||||
})
|
||||
|
||||
return queries, docs, qrels
|
||||
|
||||
|
||||
def main():
|
||||
dataset = datasets.load_dataset(args.dataset_name, cache_dir=args.cache_dir)
|
||||
|
||||
name = "swe-bench"
|
||||
if "lite" in args.dataset_name.lower():
|
||||
name += "-lite"
|
||||
|
||||
path = os.path.join(args.output_dir, name)
|
||||
os.makedirs(path, exist_ok=True)
|
||||
os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
|
||||
|
||||
queries, docs, qrels = document2code(dataset, split="test")
|
||||
save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
|
||||
save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
|
||||
qrels_path = os.path.join(path, "qrels", "test.tsv")
|
||||
save_tsv_dict(qrels, qrels_path, ["query-id", "corpus-id", "score"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--dataset_name", type=str, default="princeton-nlp/SWE-bench_Lite",
|
||||
choices=["princeton-nlp/SWE-bench", "princeton-nlp/SWE-bench_Lite"])
|
||||
parser.add_argument("--cache_dir", type=str, default="/scratch/zhiruow/data")
|
||||
parser.add_argument("--tmp_dir", type=str, default="/scratch/zhiruow/tmp")
|
||||
parser.add_argument("--output_dir", type=str, default="datasets")
|
||||
parser.add_argument("--num_examples", type=int, default=None)
|
||||
args = parser.parse_args()
|
||||
|
||||
main()
|
||||
@@ -0,0 +1,263 @@
|
||||
import os
|
||||
import re
|
||||
import chardet
|
||||
import unidiff
|
||||
import argparse
|
||||
import datasets
|
||||
import traceback
|
||||
import subprocess
|
||||
from git import Repo
|
||||
from tqdm import tqdm
|
||||
from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
from create.utils import save_tsv_dict, save_file_jsonl
|
||||
|
||||
# %% Get oracle file contents
|
||||
|
||||
# get oracle file contents from the repo
|
||||
class ContextManager:
|
||||
def __init__(self, repo_path, base_commit, verbose=False):
|
||||
self.repo_path = Path(repo_path).resolve().as_posix()
|
||||
self.old_dir = os.getcwd()
|
||||
self.base_commit = base_commit
|
||||
self.verbose = verbose
|
||||
|
||||
def __enter__(self):
|
||||
os.chdir(self.repo_path)
|
||||
cmd = f"git reset --hard {self.base_commit} && git clean -fdxq"
|
||||
if self.verbose:
|
||||
subprocess.run(cmd, shell=True, check=True)
|
||||
else:
|
||||
subprocess.run(
|
||||
cmd,
|
||||
shell=True,
|
||||
check=True,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
return self
|
||||
|
||||
def get_environment(self):
|
||||
raise NotImplementedError() # TODO: activate conda environment and return the environment file
|
||||
|
||||
def get_readme_files(self):
|
||||
files = os.listdir(self.repo_path)
|
||||
files = list(filter(lambda x: os.path.isfile(x), files))
|
||||
files = list(filter(lambda x: x.lower().startswith("readme"), files))
|
||||
return files
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
os.chdir(self.old_dir)
|
||||
|
||||
|
||||
class AutoContextManager(ContextManager):
|
||||
"""Automatically clones the repo if it doesn't exist"""
|
||||
|
||||
def __init__(self, instance, root_dir=None, verbose=False, token=None):
|
||||
if token is None:
|
||||
token = os.environ.get("GITHUB_TOKEN", "git")
|
||||
self.tempdir = None
|
||||
if root_dir is None:
|
||||
self.tempdir = TemporaryDirectory()
|
||||
root_dir = self.tempdir.name
|
||||
self.root_dir = root_dir
|
||||
repo_dir = os.path.join(self.root_dir, instance["repo"].replace("/", "__"))
|
||||
if not os.path.exists(repo_dir):
|
||||
repo_url = (
|
||||
f"https://{token}@github.com/swe-bench/"
|
||||
+ instance["repo"].replace("/", "__")
|
||||
+ ".git"
|
||||
)
|
||||
if verbose:
|
||||
print(f"Cloning {instance['repo']} to {root_dir}")
|
||||
Repo.clone_from(repo_url, repo_dir)
|
||||
super().__init__(repo_dir, instance["base_commit"], verbose=verbose)
|
||||
self.instance = instance
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
if self.tempdir is not None:
|
||||
self.tempdir.cleanup()
|
||||
return super().__exit__(exc_type, exc_val, exc_tb)
|
||||
|
||||
|
||||
def ingest_files(filenames):
|
||||
files_dict = dict()
|
||||
for filename in filenames:
|
||||
with open(filename) as f:
|
||||
content = f.read()
|
||||
files_dict[filename] = content
|
||||
return files_dict
|
||||
|
||||
def get_oracle_filenames(instance):
|
||||
"""
|
||||
Returns the filenames that are changed in the patch
|
||||
"""
|
||||
source_files = {
|
||||
patch_file.source_file.split("a/", 1)[-1]
|
||||
for patch_file in unidiff.PatchSet(instance["patch"])
|
||||
}
|
||||
gold_docs = set()
|
||||
for source_file in source_files:
|
||||
gold_docs.add(source_file)
|
||||
return gold_docs
|
||||
|
||||
|
||||
# get all file contents from the repo
|
||||
def is_test(name, test_phrases=None):
|
||||
if test_phrases is None:
|
||||
test_phrases = ["test", "tests", "testing"]
|
||||
words = set(re.split(r" |_|\/|\.", name.lower()))
|
||||
return any(word in words for word in test_phrases)
|
||||
|
||||
def list_files(root_dir, include_tests=False):
|
||||
files = []
|
||||
for filename in Path(root_dir).rglob("*.py"):
|
||||
if not include_tests and is_test(filename.as_posix()):
|
||||
continue
|
||||
files.append(filename.relative_to(root_dir).as_posix())
|
||||
return files
|
||||
|
||||
def detect_encoding(filename):
|
||||
"""
|
||||
Detect the encoding of a file
|
||||
"""
|
||||
with open(filename, "rb") as file:
|
||||
rawdata = file.read()
|
||||
return chardet.detect(rawdata)["encoding"]
|
||||
|
||||
def ingest_directory_contents(root_dir, include_tests=False):
|
||||
files_content = {}
|
||||
for relative_path in list_files(root_dir, include_tests=include_tests):
|
||||
filename = os.path.join(root_dir, relative_path)
|
||||
encoding = detect_encoding(filename)
|
||||
if encoding is None:
|
||||
content = "[BINARY DATA FILE]"
|
||||
else:
|
||||
try:
|
||||
with open(filename, encoding=encoding) as file:
|
||||
content = file.read()
|
||||
except (UnicodeDecodeError, LookupError):
|
||||
content = "[BINARY DATA FILE]"
|
||||
files_content[relative_path] = content
|
||||
return files_content
|
||||
|
||||
def get_file_contents(input_instances, verbose: bool = False, tmp_dir: str = "/scratch"):
|
||||
orig_dir = os.getcwd()
|
||||
with TemporaryDirectory(dir=tmp_dir if os.path.exists(tmp_dir) else "/tmp") as root_dir:
|
||||
for instance_id, instance in tqdm(
|
||||
input_instances.items(),
|
||||
total=len(input_instances),
|
||||
desc="Getting file contents",
|
||||
):
|
||||
try:
|
||||
with AutoContextManager(instance, root_dir, verbose=verbose) as cm:
|
||||
readmes = cm.get_readme_files()
|
||||
instance["readmes"] = ingest_files(readmes)
|
||||
instance["oracle_file_contents"] = ingest_files(get_oracle_filenames(instance))
|
||||
instance["file_contents"] = ingest_directory_contents(cm.repo_path)
|
||||
assert all([
|
||||
okey in instance["file_contents"]
|
||||
for okey in instance["oracle_file_contents"].keys()
|
||||
])
|
||||
except Exception as e:
|
||||
print(f"Failed on instance {instance_id}", e)
|
||||
traceback.print_exc()
|
||||
finally:
|
||||
# if AutoContextManager fails to exit properly future exits will return the wrong directory
|
||||
os.chdir(orig_dir)
|
||||
os.chdir(orig_dir)
|
||||
|
||||
|
||||
import multiprocessing as mp
|
||||
from functools import partial
|
||||
|
||||
def process_single_item(item, args):
|
||||
"""处理单个数据项的函数"""
|
||||
name = "swe-bench"
|
||||
if "lite" in args.dataset_name.lower():
|
||||
name += "-lite"
|
||||
|
||||
queries = [{
|
||||
"_id": item["instance_id"],
|
||||
"text": item["problem_statement"],
|
||||
"metadata": {}
|
||||
}]
|
||||
item_dict = {item["instance_id"]: item}
|
||||
|
||||
output_path = os.path.join(args.output_dir, f"{name}_{item['instance_id']}", "qrels", "test.tsv")
|
||||
if os.path.exists(output_path):
|
||||
return
|
||||
|
||||
try:
|
||||
get_file_contents(item_dict, tmp_dir=args.tmp_dir)
|
||||
|
||||
docs = []
|
||||
for instance_id, instance in item_dict.items():
|
||||
print(f"Instance #{instance_id}: {len(instance['oracle_file_contents'])} oracle / {len(instance['file_contents'])} files")
|
||||
for filename, content in instance["file_contents"].items():
|
||||
docs.append({
|
||||
"_id": f"{instance_id}_{filename}",
|
||||
"title": filename,
|
||||
"text": content,
|
||||
"metadata": {},
|
||||
})
|
||||
|
||||
qrels = []
|
||||
for instance_id, instance in item_dict.items():
|
||||
for filename, content in instance["oracle_file_contents"].items():
|
||||
qrels.append({
|
||||
"query-id": instance_id,
|
||||
"corpus-id": f"{instance_id}_{filename}",
|
||||
"score": 1
|
||||
})
|
||||
|
||||
path = os.path.join(args.output_dir, f"{name}_{instance_id}")
|
||||
os.makedirs(path, exist_ok=True)
|
||||
os.makedirs(os.path.join(path, "qrels"), exist_ok=True)
|
||||
|
||||
save_file_jsonl(queries, os.path.join(path, "queries.jsonl"))
|
||||
save_file_jsonl(docs, os.path.join(path, "corpus.jsonl"))
|
||||
qrels_path = os.path.join(path, "qrels", "test.tsv")
|
||||
save_tsv_dict(qrels, qrels_path, ["query-id", "corpus-id", "score"])
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing item {item['instance_id']}: {str(e)}")
|
||||
|
||||
def main():
|
||||
dataset = datasets.load_dataset(args.dataset_name, cache_dir=args.cache_dir)["test"]
|
||||
if args.num_examples is not None:
|
||||
import random
|
||||
indices = random.sample([i for i in range(len(dataset))], args.num_examples)
|
||||
dataset = dataset.select(indices)
|
||||
print(dataset)
|
||||
|
||||
# 创建进程池
|
||||
num_processes = mp.cpu_count() - 1 # 留一个CPU核心
|
||||
pool = mp.Pool(processes=num_processes)
|
||||
|
||||
# 使用partial固定args参数
|
||||
process_func = partial(process_single_item, args=args)
|
||||
|
||||
# 使用进程池并行处理
|
||||
list(tqdm(
|
||||
pool.imap_unordered(process_func, dataset),
|
||||
total=len(dataset),
|
||||
desc="Processing items"
|
||||
))
|
||||
|
||||
# 关闭进程池
|
||||
pool.close()
|
||||
pool.join()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--dataset_name", type=str, default="princeton-nlp/SWE-bench_Lite",
|
||||
choices=["princeton-nlp/SWE-bench", "princeton-nlp/SWE-bench_Lite"])
|
||||
parser.add_argument("--cache_dir", type=str, default="/scratch/zhiruow/data")
|
||||
parser.add_argument("--tmp_dir", type=str, default="/scratch/zhiruow/tmp")
|
||||
parser.add_argument("--output_dir", type=str, default="datasets")
|
||||
parser.add_argument("--num_examples", type=int, default=None)
|
||||
args = parser.parse_args()
|
||||
|
||||
main()
|
||||
@@ -0,0 +1,37 @@
|
||||
import jsonlines
|
||||
import csv
|
||||
import os
|
||||
|
||||
def load_jsonlines(file):
|
||||
with jsonlines.open(file, 'r') as jsonl_f:
|
||||
lst = [obj for obj in jsonl_f]
|
||||
return lst
|
||||
|
||||
def save_file_jsonl(data, fp):
|
||||
with jsonlines.open(fp, mode='w') as writer:
|
||||
writer.write_all(data)
|
||||
|
||||
def save_tsv_dict(data, fp, fields):
|
||||
# build dir
|
||||
dir_path = os.path.dirname(fp)
|
||||
os.makedirs(dir_path, exist_ok=True)
|
||||
|
||||
# writing to csv file
|
||||
with open(fp, 'w') as csvfile:
|
||||
writer = csv.DictWriter(csvfile, fieldnames=fields, delimiter='\t',)
|
||||
writer.writeheader()
|
||||
writer.writerows(data)
|
||||
|
||||
def cost_esitmate(path):
|
||||
corpus = load_jsonlines(os.path.join(path, "corpus.jsonl"))
|
||||
queries = load_jsonlines(os.path.join(path, "queries.jsonl"))
|
||||
num_corpus_words = 0
|
||||
num_queries_words = 0
|
||||
for item in tqdm(corpus):
|
||||
num_corpus_words += len(item["text"].split(" "))
|
||||
for item in tqdm(queries):
|
||||
num_queries_words += len(item["text"].split(" "))
|
||||
print(len(corpus))
|
||||
print(len(queries))
|
||||
print(num_corpus_words)
|
||||
print(num_queries_words)
|
||||
@@ -0,0 +1,313 @@
|
||||
import os
|
||||
import json
|
||||
import random
|
||||
import logging
|
||||
import pathlib
|
||||
import argparse
|
||||
import numpy as np
|
||||
from time import time
|
||||
from datasets import load_dataset
|
||||
from beir import util, LoggingHandler
|
||||
from beir.retrieval import models
|
||||
from beir.datasets.data_loader import GenericDataLoader
|
||||
from beir.retrieval.evaluation import EvaluateRetrieval
|
||||
from beir.retrieval.search.dense import DenseRetrievalExactSearch as DRES
|
||||
from tqdm import tqdm
|
||||
from transformers import HfArgumentParser
|
||||
|
||||
from arguments import CodeRAGEvalArgs, CodeRAGEvalModelArgs
|
||||
from prompts import get_task_def_by_task_name
|
||||
from FlagEmbedding import FlagLLMModel, FlagModel
|
||||
|
||||
|
||||
def get_model(model_args: CodeRAGEvalModelArgs):
|
||||
embedder_name_or_path = model_args.embedder_name_or_path
|
||||
|
||||
if model_args.embedder_model_class == "encoder-only-base":
|
||||
embedder = FlagModel(
|
||||
model_name_or_path=embedder_name_or_path,
|
||||
normalize_embeddings=model_args.normalize_embeddings,
|
||||
pooling_method=model_args.pooling_method,
|
||||
use_fp16=model_args.use_fp16,
|
||||
query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
|
||||
query_instruction_format=model_args.query_instruction_format_for_retrieval,
|
||||
devices=model_args.devices,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
cache_dir=model_args.cache_dir,
|
||||
batch_size=model_args.embedder_batch_size,
|
||||
query_max_length=model_args.embedder_query_max_length,
|
||||
passage_max_length=model_args.embedder_passage_max_length,
|
||||
)
|
||||
elif model_args.embedder_model_class == "decoder-only-base":
|
||||
embedder = FlagLLMModel(
|
||||
model_name_or_path=embedder_name_or_path,
|
||||
normalize_embeddings=model_args.normalize_embeddings,
|
||||
pooling_method=model_args.pooling_method,
|
||||
use_fp16=model_args.use_fp16,
|
||||
query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
|
||||
query_instruction_format=model_args.query_instruction_format_for_retrieval,
|
||||
devices=model_args.devices,
|
||||
examples_for_task=model_args.examples_for_task,
|
||||
examples_instruction_format=model_args.examples_instruction_format,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
cache_dir=model_args.cache_dir,
|
||||
batch_size=model_args.embedder_batch_size,
|
||||
query_max_length=model_args.embedder_query_max_length,
|
||||
passage_max_length=model_args.embedder_passage_max_length,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid model class: {model_args.embedder_model_class}")
|
||||
embedder.model.config._name_or_path = model_args.embedder_name_or_path
|
||||
|
||||
class CustomFlagModel:
|
||||
def __init__(self, model):
|
||||
self.model = model
|
||||
|
||||
def encode_queries(self, queries, show_progress_bar, convert_to_tensor, **kwargs):
|
||||
if isinstance(queries, str):
|
||||
queries = [queries]
|
||||
|
||||
if isinstance(queries[0], dict):
|
||||
queries = [(e.get('title') + ' ' + e['text']).strip() for e in queries]
|
||||
|
||||
return self.model.encode_queries(queries, **kwargs)
|
||||
|
||||
def encode_corpus(self, corpus, show_progress_bar, convert_to_tensor, **kwargs):
|
||||
if isinstance(corpus, str):
|
||||
corpus = [corpus]
|
||||
|
||||
if isinstance(corpus[0], dict):
|
||||
corpus = [(e.get('title') + ' ' + e['text']).strip() for e in corpus]
|
||||
|
||||
return self.model.encode_corpus(corpus, **kwargs)
|
||||
|
||||
def encode(self, corpus, show_progress_bar, convert_to_tensor, **kwargs):
|
||||
if isinstance(corpus, str):
|
||||
corpus = [corpus]
|
||||
|
||||
if isinstance(corpus[0], dict):
|
||||
corpus = [(e.get('title') + ' ' + e['text']).strip() for e in corpus]
|
||||
|
||||
return self.model.encode(corpus, **kwargs)
|
||||
|
||||
return CustomFlagModel(embedder)
|
||||
|
||||
#### Just some code to print debug information to stdout
|
||||
logging.basicConfig(format='%(asctime)s - %(message)s',
|
||||
datefmt='%Y-%m-%d %H:%M:%S',
|
||||
level=logging.INFO,
|
||||
handlers=[LoggingHandler()])
|
||||
|
||||
|
||||
def get_top_docs(results: dict, corpus: dict, task_id: str, topk: int = 10) -> list[str]:
|
||||
if task_id not in results: return []
|
||||
doc_scores = results[task_id]
|
||||
doc_scores_sorted = sorted(doc_scores.items(), key=lambda item: item[1], reverse=True)
|
||||
doc_scores_sorted = doc_scores_sorted[:topk]
|
||||
doc_code_snippets = [corpus[code_id] for code_id, score in doc_scores_sorted]
|
||||
return doc_code_snippets
|
||||
|
||||
|
||||
def main(
|
||||
eval_args: CodeRAGEvalArgs,
|
||||
model_args: CodeRAGEvalModelArgs
|
||||
):
|
||||
args = eval_args
|
||||
|
||||
embedder = get_model(model_args)
|
||||
model = DRES(
|
||||
embedder,
|
||||
batch_size=args.batch_size,
|
||||
corpus_chunk_size=512 * 9999
|
||||
)
|
||||
retriever = EvaluateRetrieval(model, score_function="dot")
|
||||
|
||||
if args.dataset.startswith("swe-bench") or args.dataset.startswith("repoeval"):
|
||||
all_eval_results = []
|
||||
|
||||
if args.dataset.startswith("swe-bench"):
|
||||
swebench = load_dataset("princeton-nlp/SWE-bench_Lite")["test"]
|
||||
all_top_docs = [[] for _ in swebench]
|
||||
|
||||
instance_list = [i for i in os.listdir("datasets") if i.startswith(f"{args.dataset}_")]
|
||||
instance_list_filtered = []
|
||||
|
||||
for ins_dir in tqdm(instance_list):
|
||||
logging.info("Instance Repo: {}".format(ins_dir))
|
||||
# load data and perform retrieval
|
||||
corpus, queries, qrels = GenericDataLoader(
|
||||
data_folder=os.path.join("datasets", ins_dir)
|
||||
).load(split="test")
|
||||
logging.info(f"Instance #{ins_dir}: #{len(corpus)} corpus, #{len(queries)} queries")
|
||||
|
||||
start_time = time()
|
||||
if len(queries) == 1:
|
||||
queries.update({"dummy": "dummy"})
|
||||
results = retriever.retrieve(corpus, queries)
|
||||
if "dummy" in queries:
|
||||
queries.pop("dummy")
|
||||
results.pop("dummy")
|
||||
end_time = time()
|
||||
logging.info("Time taken to retrieve: {:.2f} seconds".format(end_time - start_time))
|
||||
|
||||
# get topk retrieved docs
|
||||
if args.dataset.startswith("swe-bench"):
|
||||
indices = [i for i, ex in enumerate(swebench) if ex["instance_id"] in queries]
|
||||
for index in indices:
|
||||
instance_id = swebench[index]["instance_id"]
|
||||
all_top_docs[index] = get_top_docs(results, corpus, instance_id)
|
||||
elif args.dataset.startswith("repoeval"):
|
||||
args.dataset_path = "output/repoeval/datasets/function_level_completion_2k_context_codex.test.clean.jsonl"
|
||||
tasks = [json.loads(line.strip()) for line in open(args.dataset_path, 'r')]
|
||||
prompts, references, docs, metadatas = [], [], [], []
|
||||
for task in tasks:
|
||||
if task["metadata"]["task_id"] not in queries: continue
|
||||
prompts.append(task["prompt"]) # save full prompt
|
||||
references.append(task["metadata"]["ground_truth"])
|
||||
docs.append(get_top_docs(
|
||||
results=results, corpus=corpus, task_id=task["metadata"]["task_id"],
|
||||
))
|
||||
metadatas.append(task["metadata"])
|
||||
assert len(prompts) == len(references) == len(docs)
|
||||
dataset = [
|
||||
{"prompt": p, "reference": r, "docs": d, "metadata": m}
|
||||
for p, r, d, m in zip(prompts, references, docs, metadatas)
|
||||
]
|
||||
with open(args.results_file, "a") as fout:
|
||||
for curr in dataset:
|
||||
fout.write(json.dumps(curr) + "\n")
|
||||
else:
|
||||
raise ValueError(f"`dataset` should starts with either 'swe-bench' or 'repoeval'.")
|
||||
|
||||
# evaluate retrieval results
|
||||
if len(qrels) == 0:
|
||||
logging.info("No qrels found for this dataset.")
|
||||
return
|
||||
logging.info("Retriever evaluation for k in: {}".format(retriever.k_values))
|
||||
ndcg, _map, recall, precision = retriever.evaluate(qrels, results, retriever.k_values)
|
||||
mrr = retriever.evaluate_custom(qrels, results, retriever.k_values, metric="mrr")
|
||||
eval_results = {
|
||||
"ndcg": ndcg, "mrr": mrr,
|
||||
"recall": recall, "precision": precision,
|
||||
"time": end_time - start_time
|
||||
}
|
||||
logging.info(f"Instance #{ins_dir}: {eval_results}")
|
||||
all_eval_results.append(eval_results)
|
||||
|
||||
with open(args.output_file + "_all", "w") as f:
|
||||
json.dump(all_eval_results, f)
|
||||
|
||||
if args.dataset.startswith("swe-bench"):
|
||||
swebench = swebench.add_column("docs", all_top_docs)
|
||||
swebench.to_json(args.results_file)
|
||||
|
||||
avg_eval_results = {}
|
||||
for k, v_dict in all_eval_results[0].items():
|
||||
if isinstance(v_dict, dict):
|
||||
avg_v_dict = {}
|
||||
for vk, vv in v_dict.items():
|
||||
avg_vv = sum([e[k][vk] for e in all_eval_results]) / len(all_eval_results)
|
||||
avg_v_dict[vk] = avg_vv
|
||||
avg_eval_results.update(avg_v_dict)
|
||||
elif isinstance(v_dict, float):
|
||||
avg_v = sum([e[k] for e in all_eval_results]) / len(all_eval_results)
|
||||
avg_eval_results[k] = avg_v
|
||||
else:
|
||||
raise ValueError
|
||||
print("Average Eval Results: ", avg_eval_results)
|
||||
with open(args.output_file, "w") as f:
|
||||
json.dump(avg_eval_results, f)
|
||||
else:
|
||||
dataset = args.dataset
|
||||
corpus, queries, qrels = GenericDataLoader(data_folder=os.path.join("datasets", args.dataset)).load(
|
||||
split="test")
|
||||
#### Retrieve dense results (format of results is identical to qrels)
|
||||
start_time = time()
|
||||
results = retriever.retrieve(corpus, queries)
|
||||
end_time = time()
|
||||
print("Time taken to retrieve: {:.2f} seconds".format(end_time - start_time))
|
||||
|
||||
if args.dataset in ["humaneval", "mbpp", "apps"]:
|
||||
if args.dataset == "humaneval":
|
||||
ds = load_dataset("openai_humaneval")
|
||||
id_key = "task_id"
|
||||
elif args.dataset == "mbpp":
|
||||
ds = load_dataset("mbpp")
|
||||
id_key = "task_id"
|
||||
elif args.dataset == "apps":
|
||||
ds = load_dataset("codeparrot/apps")
|
||||
id_key = "problem_id"
|
||||
all_top_docs = []
|
||||
for task_id in ds["test"][id_key]:
|
||||
all_top_docs.append(get_top_docs(results, corpus, f"{task_id}_doc"))
|
||||
ds["test"] = ds["test"].add_column("docs", all_top_docs)
|
||||
ds["test"].to_json(args.results_file) # this outputs to arrow format and read as .jsonl
|
||||
elif args.dataset.startswith("odex"):
|
||||
lang = args.dataset.split("_")[-1]
|
||||
ds = load_dataset("neulab/odex", lang, trust_remote_code=True)
|
||||
all_top_docs = []
|
||||
for idx, task_id in enumerate(ds["test"]["task_id"]):
|
||||
all_top_docs.append(get_top_docs(results, corpus, f"{idx}_{task_id}"))
|
||||
ds["test"] = ds["test"].add_column("docs", all_top_docs)
|
||||
ds["test"].to_json(args.results_file) # this outputs to arrow format and read as .jsonl
|
||||
elif args.dataset.startswith("ds1000"):
|
||||
_, key, mode = args.dataset.split("_")
|
||||
key = key.capitalize()
|
||||
mode = mode.capitalize()
|
||||
from create.ds1000 import get_dataset
|
||||
source_dir = pathlib.Path(__file__).parent / "ds"
|
||||
data = get_dataset(source_dir, mode=mode, key=key)
|
||||
all_docs = []
|
||||
example_ids = []
|
||||
for item in data:
|
||||
example = item.data
|
||||
example_id = f"{example['lib']}_{example['perturbation_origin_id']}"
|
||||
all_docs.append(get_top_docs(results, corpus, example_id))
|
||||
example_ids.append(example_id)
|
||||
assert len(all_docs) == len(
|
||||
example_ids), f"length of all_docs should be {len(example_ids)}, now is {len(all_docs)}"
|
||||
with open(args.results_file, "w+") as fout:
|
||||
for idx, all_doc in enumerate(all_docs):
|
||||
fout.write(json.dumps({"example_id": example_id,
|
||||
"docs": all_doc}) + "\n")
|
||||
else:
|
||||
with open(args.results_file, 'w+') as fw:
|
||||
for curr in results:
|
||||
fw.write(json.dumps({curr: results[curr]}) + "\n")
|
||||
|
||||
#### Evaluate your retrieval using NDCG@k, MAP@K ...
|
||||
if len(qrels) == 0:
|
||||
logging.info("No qrels found for this dataset.")
|
||||
return
|
||||
logging.info("Retriever evaluation for k in: {}".format(retriever.k_values))
|
||||
ndcg, _map, recall, precision = retriever.evaluate(qrels, results, retriever.k_values)
|
||||
|
||||
mrr = retriever.evaluate_custom(qrels, results, retriever.k_values, metric="mrr")
|
||||
recall_cap = retriever.evaluate_custom(qrels, results, retriever.k_values, metric="r_cap")
|
||||
hole = retriever.evaluate_custom(qrels, results, retriever.k_values, metric="hole")
|
||||
|
||||
all_results = {"ndcg": ndcg, "mrr": mrr, "recall": recall, "precision": precision,
|
||||
"time": end_time - start_time}
|
||||
with open(args.output_file, "w") as f:
|
||||
json.dump(all_results, f)
|
||||
#### Print top-k documents retrieved ####
|
||||
top_k = 3
|
||||
|
||||
query_id, ranking_scores = random.choice(list(results.items()))
|
||||
scores_sorted = sorted(ranking_scores.items(), key=lambda item: item[1], reverse=True)
|
||||
logging.info("Query : %s\n" % queries[query_id])
|
||||
|
||||
for rank in range(top_k):
|
||||
doc_id = scores_sorted[rank][0]
|
||||
# Format: Rank x: ID [Title] Body
|
||||
logging.info(
|
||||
"Rank %d: %s [%s] - %s\n" % (rank + 1, doc_id, corpus[doc_id].get("title"), corpus[doc_id].get("text")))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = HfArgumentParser((
|
||||
CodeRAGEvalArgs,
|
||||
CodeRAGEvalModelArgs
|
||||
))
|
||||
eval_args, model_args = parser.parse_args_into_dataclasses()
|
||||
main(eval_args, model_args)
|
||||
@@ -0,0 +1,18 @@
|
||||
from typing import Dict
|
||||
|
||||
|
||||
def get_task_def_by_task_name(task_name: str) -> str:
|
||||
task_name_to_instruct: Dict[str, str] = {
|
||||
'humaneval': 'Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
|
||||
'mbpp': 'Given a textual explanation of code functionality, retrieve the corresponding code implementation.',
|
||||
'ds1000_all_completion': 'Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
|
||||
'odex_en': 'Given a question, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
|
||||
'odex_es': 'Given a question, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
|
||||
'odex_ja': 'Given a question, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
|
||||
'odex_ru': 'Given a question, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
|
||||
'repoeval': 'Given a code snippet and a new function name, retrieve the implementation of the function.',
|
||||
# 'repoeval': 'Given a piece of code segment, retrieve the code segment that is the latter part of the code.',
|
||||
'swe-bench-lite': 'Given a code snippet containing a bug and a natural language description of the bug or error, retrieve code snippets that demonstrate solutions or fixes for similar bugs or errors (the desired documents).'
|
||||
}
|
||||
|
||||
return task_name_to_instruct[task_name]
|
||||
@@ -0,0 +1,69 @@
|
||||
from typing import List
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from FlagEmbedding.abc.evaluation import (
|
||||
AbsEvalModelArgs as COIREvalModelArgs,
|
||||
)
|
||||
|
||||
|
||||
def coir_tasks():
|
||||
return [
|
||||
"apps",
|
||||
"codefeedback-mt",
|
||||
"codefeedback-st",
|
||||
"CodeSearchNet-ccr-go",
|
||||
"CodeSearchNet-ccr-java",
|
||||
"CodeSearchNet-ccr-javascript",
|
||||
"CodeSearchNet-ccr-php",
|
||||
"CodeSearchNet-ccr-python",
|
||||
"CodeSearchNet-ccr-ruby",
|
||||
"CodeSearchNet-go",
|
||||
"CodeSearchNet-java",
|
||||
"CodeSearchNet-javascript",
|
||||
"CodeSearchNet-php",
|
||||
"CodeSearchNet-python",
|
||||
"CodeSearchNet-ruby",
|
||||
"codetrans-contest",
|
||||
"codetrans-dl",
|
||||
"cosqa",
|
||||
"stackoverflow-qa",
|
||||
"synthetic-text2sql"
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class COIREvalArgs:
|
||||
output_dir: str = field(
|
||||
default="./results", metadata={"help": "Path to save results."}
|
||||
)
|
||||
tasks: List[str] = field(
|
||||
default_factory=coir_tasks,
|
||||
metadata={
|
||||
"help": "Tasks to evaluate. Default: None. Available tasks: ['apps', 'codefeedback-mt', 'codefeedback-st', 'CodeSearchNet-ccr-go', 'CodeSearchNet-ccr-java', 'CodeSearchNet-ccr-javascript', 'CodeSearchNet-ccr-php', 'CodeSearchNet-ccr-python', 'CodeSearchNet-ccr-ruby', 'CodeSearchNet-go', 'CodeSearchNet-java', 'CodeSearchNet-javascript', 'CodeSearchNet-php', 'CodeSearchNet-python', 'CodeSearchNet-ruby', 'codetrans-contest', 'codetrans-dl', 'cosqa', 'stackoverflow-qa', 'synthetic-text2sql']",
|
||||
"choices": [
|
||||
"apps",
|
||||
"codefeedback-mt",
|
||||
"codefeedback-st",
|
||||
"CodeSearchNet-ccr-go",
|
||||
"CodeSearchNet-ccr-java",
|
||||
"CodeSearchNet-ccr-javascript",
|
||||
"CodeSearchNet-ccr-php",
|
||||
"CodeSearchNet-ccr-python",
|
||||
"CodeSearchNet-ccr-ruby",
|
||||
"CodeSearchNet-go",
|
||||
"CodeSearchNet-java",
|
||||
"CodeSearchNet-javascript",
|
||||
"CodeSearchNet-php",
|
||||
"CodeSearchNet-python",
|
||||
"CodeSearchNet-ruby",
|
||||
"codetrans-contest",
|
||||
"codetrans-dl",
|
||||
"cosqa",
|
||||
"stackoverflow-qa",
|
||||
"synthetic-text2sql"
|
||||
]
|
||||
}
|
||||
)
|
||||
use_special_instructions: bool = field(
|
||||
default=False, metadata={"help": "Whether to use specific instructions in `prompts.py` for evaluation. Default: False"}
|
||||
)
|
||||
@@ -0,0 +1,16 @@
|
||||
output_dir=result
|
||||
|
||||
python main.py \
|
||||
--output_dir ${output_dir} \
|
||||
--use_special_instructions True \
|
||||
--embedder_name_or_path BAAI/bge-code-v1 \
|
||||
--embedder_model_class decoder-only-base \
|
||||
--query_instruction_format_for_retrieval '<instruct>{}\n<query>{}' \
|
||||
--embedder_query_max_length 2048 \
|
||||
--embedder_passage_max_length 2048 \
|
||||
--trust_remote_code True \
|
||||
--pooling_method last_token \
|
||||
--embedder_batch_size 64 \
|
||||
--devices cuda:0 cuda:1 cuda:2 cuda:3 cuda:4 cuda:5 cuda:6 cuda:7 \
|
||||
--tasks apps codetrans-contest codetrans-dl cosqa synthetic-text2sql stackoverflow-qa codefeedback-mt codefeedback-st CodeSearchNet-ccr-go CodeSearchNet-ccr-java CodeSearchNet-ccr-javascript CodeSearchNet-ccr-php CodeSearchNet-ccr-python CodeSearchNet-ccr-ruby CodeSearchNet-go CodeSearchNet-java CodeSearchNet-javascript CodeSearchNet-php CodeSearchNet-python CodeSearchNet-ruby \
|
||||
--cache_dir ./cache
|
||||
@@ -0,0 +1,188 @@
|
||||
import os
|
||||
import json
|
||||
import coir
|
||||
from transformers import HfArgumentParser
|
||||
|
||||
from arguments import COIREvalArgs, COIREvalModelArgs
|
||||
from prompts import get_task_def_by_task_name
|
||||
from FlagEmbedding import FlagLLMModel, FlagModel, FlagPseudoMoEModel
|
||||
|
||||
|
||||
def get_model(model_args: COIREvalModelArgs):
|
||||
embedder_name_or_path = model_args.embedder_name_or_path
|
||||
|
||||
if model_args.embedder_model_class == "encoder-only-base":
|
||||
embedder = FlagModel(
|
||||
model_name_or_path=embedder_name_or_path,
|
||||
normalize_embeddings=model_args.normalize_embeddings,
|
||||
pooling_method=model_args.pooling_method,
|
||||
use_fp16=model_args.use_fp16,
|
||||
use_bf16=model_args.use_bf16,
|
||||
query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
|
||||
query_instruction_format=model_args.query_instruction_format_for_retrieval,
|
||||
devices=model_args.devices,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
cache_dir=model_args.cache_dir,
|
||||
batch_size=model_args.embedder_batch_size,
|
||||
query_max_length=model_args.embedder_query_max_length,
|
||||
passage_max_length=model_args.embedder_passage_max_length,
|
||||
)
|
||||
elif model_args.embedder_model_class == "decoder-only-base":
|
||||
embedder = FlagLLMModel(
|
||||
model_name_or_path=embedder_name_or_path,
|
||||
normalize_embeddings=model_args.normalize_embeddings,
|
||||
pooling_method=model_args.pooling_method,
|
||||
use_fp16=model_args.use_fp16,
|
||||
use_bf16=model_args.use_bf16,
|
||||
query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
|
||||
query_instruction_format=model_args.query_instruction_format_for_retrieval,
|
||||
devices=model_args.devices,
|
||||
examples_for_task=model_args.examples_for_task,
|
||||
examples_instruction_format=model_args.examples_instruction_format,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
cache_dir=model_args.cache_dir,
|
||||
batch_size=model_args.embedder_batch_size,
|
||||
query_max_length=model_args.embedder_query_max_length,
|
||||
passage_max_length=model_args.embedder_passage_max_length,
|
||||
)
|
||||
elif model_args.embedder_model_class == "decoder-only-pseudo_moe":
|
||||
embedder = FlagPseudoMoEModel(
|
||||
model_name_or_path=embedder_name_or_path,
|
||||
normalize_embeddings=model_args.normalize_embeddings,
|
||||
pooling_method=model_args.pooling_method,
|
||||
use_fp16=model_args.use_fp16,
|
||||
use_bf16=model_args.use_bf16,
|
||||
query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
|
||||
query_instruction_format=model_args.query_instruction_format_for_retrieval,
|
||||
devices=model_args.devices,
|
||||
examples_for_task=model_args.examples_for_task,
|
||||
examples_instruction_format=model_args.examples_instruction_format,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
cache_dir=model_args.cache_dir,
|
||||
batch_size=model_args.embedder_batch_size,
|
||||
query_max_length=model_args.embedder_query_max_length,
|
||||
passage_max_length=model_args.embedder_passage_max_length,
|
||||
domain_for_pseudo_moe=model_args.domain_for_pseudo_moe,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid model class: {model_args.embedder_model_class}")
|
||||
embedder.model.config._name_or_path = model_args.embedder_name_or_path
|
||||
|
||||
class CustomFlagModel:
|
||||
def __init__(self, model):
|
||||
self.model = model
|
||||
|
||||
def encode_queries(self, queries, show_progress_bar, convert_to_tensor, **kwargs):
|
||||
if isinstance(queries, str):
|
||||
queries = [queries]
|
||||
|
||||
if isinstance(queries[0], dict):
|
||||
queries = [(e.get('title') + ' ' + e['text']).strip() for e in queries]
|
||||
|
||||
return self.model.encode_queries(queries, **kwargs)
|
||||
|
||||
def encode_corpus(self, corpus, show_progress_bar, convert_to_tensor, **kwargs):
|
||||
if isinstance(corpus, str):
|
||||
corpus = [corpus]
|
||||
|
||||
if isinstance(corpus[0], dict):
|
||||
corpus = [(e.get('title') + ' ' + e['text']).strip() for e in corpus]
|
||||
|
||||
return self.model.encode_corpus(corpus, **kwargs)
|
||||
|
||||
def encode(self, corpus, show_progress_bar, convert_to_tensor, **kwargs):
|
||||
if isinstance(corpus, str):
|
||||
corpus = [corpus]
|
||||
|
||||
if isinstance(corpus[0], dict):
|
||||
corpus = [(e.get('title') + ' ' + e['text']).strip() for e in corpus]
|
||||
|
||||
return self.model.encode(corpus, **kwargs)
|
||||
|
||||
return CustomFlagModel(embedder)
|
||||
|
||||
|
||||
def main(
|
||||
eval_args: COIREvalArgs,
|
||||
model_args: COIREvalModelArgs
|
||||
):
|
||||
model = get_model(model_args)
|
||||
|
||||
output_folder = os.path.join(eval_args.output_dir, os.path.basename(model.model.model.config._name_or_path))
|
||||
|
||||
all_task = eval_args.tasks
|
||||
if not isinstance(all_task, list):
|
||||
all_task = [all_task]
|
||||
|
||||
all_results = {}
|
||||
for task_name in all_task:
|
||||
save_path = os.path.join(output_folder, f"{task_name}.json")
|
||||
if os.path.exists(save_path):
|
||||
with open(save_path, "r", encoding="utf-8") as f:
|
||||
results = json.load(f)
|
||||
all_results[task_name] = results['metrics']
|
||||
continue
|
||||
|
||||
tmp_task = coir.get_tasks(tasks=[task_name])
|
||||
evaluation = coir.COIR(tasks=tmp_task,
|
||||
batch_size=model_args.embedder_batch_size)
|
||||
|
||||
model.model.stop_self_pool()
|
||||
|
||||
if eval_args.use_special_instructions:
|
||||
model.model.query_instruction_for_retrieval = get_task_def_by_task_name(task_name)
|
||||
|
||||
results = evaluation.run(model, output_folder=output_folder)
|
||||
all_results[task_name] = results[task_name]
|
||||
|
||||
csn_result = 0
|
||||
csn_num = 0
|
||||
csn_ccr_result = 0
|
||||
csn_ccr_num = 0
|
||||
pop_keys = []
|
||||
all_result = 0
|
||||
all_num = 0
|
||||
for k in all_results.keys():
|
||||
if 'CodeSearchNet-ccr' in k:
|
||||
csn_ccr_result += all_results[k]['NDCG']['NDCG@10']
|
||||
csn_ccr_num += 1
|
||||
pop_keys.append(k)
|
||||
elif 'CodeSearchNet' in k:
|
||||
csn_result += all_results[k]['NDCG']['NDCG@10']
|
||||
csn_num += 1
|
||||
pop_keys.append(k)
|
||||
else:
|
||||
all_result += all_results[k]['NDCG']['NDCG@10']
|
||||
all_num += 1
|
||||
if csn_num > 0:
|
||||
print('Using CodeSearchNet')
|
||||
all_result += csn_result / csn_num
|
||||
all_num += 1
|
||||
if csn_ccr_num > 0:
|
||||
print('Using CodeSearchNet-ccr')
|
||||
all_result += csn_ccr_result / csn_ccr_num
|
||||
all_num += 1
|
||||
new_results = {}
|
||||
for k in all_results:
|
||||
if k in pop_keys:
|
||||
continue
|
||||
new_results[k] = all_results[k]['NDCG']['NDCG@10']
|
||||
if csn_num > 0:
|
||||
new_results['CodeSearchNet'] = csn_result / csn_num
|
||||
if csn_ccr_num > 0:
|
||||
new_results['CodeSearchNet_CCR'] = csn_ccr_result / csn_ccr_num
|
||||
new_results['all'] = all_result / all_num
|
||||
|
||||
print(new_results)
|
||||
|
||||
with open(os.path.join(output_folder, 'OVERALL-results.json'), 'w') as f:
|
||||
json.dump(new_results, f)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = HfArgumentParser((
|
||||
COIREvalArgs,
|
||||
COIREvalModelArgs
|
||||
))
|
||||
eval_args, model_args = parser.parse_args_into_dataclasses()
|
||||
main(eval_args, model_args)
|
||||
@@ -0,0 +1,38 @@
|
||||
from typing import Dict
|
||||
|
||||
|
||||
def get_task_def_by_task_name(task_name: str) -> str:
|
||||
task_name_to_instruct: Dict[str, str] = {
|
||||
# Text-to-Code Retrieval
|
||||
## Code Contest Retrieval
|
||||
'apps': 'Given a code contest problem description, retrieve relevant code that can help solve the problem.',
|
||||
## Web Query to Code Retrieval
|
||||
'cosqa': 'Given a web search query, retrieve relevant code that can help answer the query.',
|
||||
## Text-to-SQL Retrieval
|
||||
'synthetic-text2sql': 'Given a question in text, retrieve SQL queries that are appropriate responses to the question.',
|
||||
|
||||
# Code-to-Text Retrieval
|
||||
## Code Summary Retrieval
|
||||
'CodeSearchNet-': 'Given a piece of code, retrieve the document string that summarizes the code.',
|
||||
|
||||
# Code-to-Code Retrieval
|
||||
## Code Context Retrieval
|
||||
'CodeSearchNet-ccr-': 'Given a piece of code segment, retrieve the code segment that is the latter part of the code.',
|
||||
## Similar Code Retrieval
|
||||
'codetrans-dl': 'Given a piece of code, retrieve code that is semantically equivalent to the input code.',
|
||||
'codetrans-contest': 'Given a piece of Python code, retrieve C++ code that is semantically equivalent to the input code.',
|
||||
|
||||
# Hybrid Code Retrieval
|
||||
## Single-turn Code QA
|
||||
'stackoverflow-qa': 'Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
|
||||
'codefeedback-st': 'Given a question that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
|
||||
## Multi-turn Code QA
|
||||
'codefeedback-mt': 'Given a multi-turn conversation history that consists of a mix of text and code snippets, retrieve relevant answers that also consist of a mix of text and code snippets, and can help answer the question.',
|
||||
}
|
||||
|
||||
special_task_names = ['CodeSearchNet-ccr-', 'CodeSearchNet-']
|
||||
for special_task_name in special_task_names:
|
||||
if special_task_name in task_name:
|
||||
return task_name_to_instruct[special_task_name]
|
||||
|
||||
return task_name_to_instruct[task_name]
|
||||
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,298 @@
|
||||
# BGE-M3 ([paper](https://arxiv.org/pdf/2402.03216.pdf), [code](https://github.com/FlagOpen/FlagEmbedding/tree/master/research/BGE_M3))
|
||||
|
||||
In this project, we introduce BGE-M3, which is distinguished for its versatility in Multi-Functionality, Multi-Linguality, and Multi-Granularity.
|
||||
- Multi-Functionality: It can simultaneously perform the three common retrieval functionalities of embedding model: dense retrieval, multi-vector retrieval, and sparse retrieval.
|
||||
- Multi-Linguality: It can support more than 100 working languages.
|
||||
- Multi-Granularity: It is able to process inputs of different granularities, spanning from short sentences to long documents of up to 8192 tokens.
|
||||
|
||||
For more details, please refer to our paper: [BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation](https://arxiv.org/pdf/2402.03216.pdf)
|
||||
|
||||
**Some suggestions for retrieval pipeline in RAG**
|
||||
|
||||
We recommend to use following pipeline: hybrid retrieval + re-ranking.
|
||||
- Hybrid retrieval leverages the strengths of various methods, offering higher accuracy and stronger generalization capabilities.
|
||||
A classic example: using both embedding retrieval and the BM25 algorithm.
|
||||
Now, you can try to use BGE-M3, which supports both embedding and sparse retrieval.
|
||||
This allows you to obtain token weights (similar to the BM25) without any additional cost when generate dense embeddings.
|
||||
To use hybrid retrieval, you can refer to [Vespa](https://github.com/vespa-engine/pyvespa/blob/master/docs/sphinx/source/examples/mother-of-all-embedding-models-cloud.ipynb
|
||||
) and [Milvus](https://github.com/milvus-io/pymilvus/blob/master/examples/hello_hybrid_sparse_dense.py).
|
||||
|
||||
- As cross-encoder models, re-ranker demonstrates higher accuracy than bi-encoder embedding model.
|
||||
Utilizing the re-ranking model (e.g., [bge-reranker](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/inference/reranker#2-normal-reranker), [bge-reranker-v2](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/inference/reranker#3-llm-based-reranker)) after retrieval can further filter the selected text.
|
||||
|
||||
|
||||
## News:
|
||||
|
||||
- 2024/7/1: **We update the MIRACL evaluation results of BGE-M3**. To reproduce the new results, you can refer to: [bge-m3_miracl_2cr](https://huggingface.co/datasets/hanhainebula/bge-m3_miracl_2cr). We have also updated our [paper](https://arxiv.org/pdf/2402.03216) on arXiv.
|
||||
|
||||
<details>
|
||||
<summary> Details </summary>
|
||||
|
||||
> The previous test results were lower because we mistakenly removed the passages that have the same id as the query from the search results. After correcting this mistake, the overall performance of BGE-M3 on MIRACL is higher than the previous results, but the experimental conclusion remains unchanged. The other results are not affected by this mistake. To reproduce the previous lower results, you need to add the `--remove-query` parameter when using `pyserini.search.faiss` or `pyserini.search.lucene` to search the passages.
|
||||
|
||||
</details>
|
||||
- 2024/3/20: **Thanks Milvus team!** Now you can use hybrid retrieval of bge-m3 in Milvus: [pymilvus/examples
|
||||
/hello_hybrid_sparse_dense.py](https://github.com/milvus-io/pymilvus/blob/master/examples/hello_hybrid_sparse_dense.py).
|
||||
- 2024/3/8: **Thanks for the [experimental results](https://towardsdatascience.com/openai-vs-open-source-multilingual-embedding-models-e5ccb7c90f05) from @[Yannael](https://huggingface.co/Yannael). In this benchmark, BGE-M3 achieves top performance in both English and other languages, surpassing models such as OpenAI.**
|
||||
- 2024/3/2: Release unified fine-tuning [example](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/finetune/embedder#2-bge-m3) and [data](https://huggingface.co/datasets/Shitao/bge-m3-data)
|
||||
- 2024/2/6: We release the [MLDR](https://huggingface.co/datasets/Shitao/MLDR) (a long document retrieval dataset covering 13 languages) and [evaluation pipeline](https://github.com/FlagOpen/FlagEmbedding/tree/master/research/C_MTEB/MLDR).
|
||||
- 2024/2/1: **Thanks for the excellent tool from Vespa.** You can easily use multiple modes of BGE-M3 following this [notebook](https://github.com/vespa-engine/pyvespa/blob/master/docs/sphinx/source/examples/mother-of-all-embedding-models-cloud.ipynb)
|
||||
|
||||
|
||||
## Specs
|
||||
|
||||
- Model
|
||||
|
||||
| Model Name | Dimension | Sequence Length | Introduction |
|
||||
|:----:|:---:|:---:|:---:|
|
||||
| [BAAI/bge-m3](https://huggingface.co/BAAI/bge-m3) | 1024 | 8192 | multilingual; unified fine-tuning (dense, sparse, and colbert) from bge-m3-unsupervised|
|
||||
| [BAAI/bge-m3-unsupervised](https://huggingface.co/BAAI/bge-m3-unsupervised) | 1024 | 8192 | multilingual; contrastive learning from bge-m3-retromae |
|
||||
| [BAAI/bge-m3-retromae](https://huggingface.co/BAAI/bge-m3-retromae) | -- | 8192 | multilingual; extend the max_length of [xlm-roberta](https://huggingface.co/FacebookAI/xlm-roberta-large) to 8192 and further pretrained via [retromae](https://github.com/staoxiao/RetroMAE)|
|
||||
| [BAAI/bge-large-en-v1.5](https://huggingface.co/BAAI/bge-large-en-v1.5) | 1024 | 512 | English model |
|
||||
| [BAAI/bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5) | 768 | 512 | English model |
|
||||
| [BAAI/bge-small-en-v1.5](https://huggingface.co/BAAI/bge-small-en-v1.5) | 384 | 512 | English model |
|
||||
|
||||
- Data
|
||||
|
||||
| Dataset | Introduction |
|
||||
|:----------------------------------------------------------:|:-------------------------------------------------:|
|
||||
| [MLDR](https://huggingface.co/datasets/Shitao/MLDR) | Docuemtn Retrieval Dataset, covering 13 languages |
|
||||
| [bge-m3-data](https://huggingface.co/datasets/Shitao/bge-m3-data) | Fine-tuning data used by bge-m3 |
|
||||
|
||||
|
||||
|
||||
## FAQ
|
||||
|
||||
**1. Introduction for different retrieval methods**
|
||||
|
||||
- Dense retrieval: map the text into a single embedding, e.g., [DPR](https://arxiv.org/abs/2004.04906), [BGE-v1.5](https://github.com/FlagOpen/FlagEmbedding)
|
||||
- Sparse retrieval (lexical matching): a vector of size equal to the vocabulary, with the majority of positions set to zero, calculating a weight only for tokens present in the text. e.g., BM25, [unicoil](https://arxiv.org/pdf/2106.14807.pdf), and [splade](https://arxiv.org/abs/2107.05720)
|
||||
- Multi-vector retrieval: use multiple vectors to represent a text, e.g., [ColBERT](https://arxiv.org/abs/2004.12832).
|
||||
|
||||
|
||||
**2. How to use BGE-M3 in other projects?**
|
||||
|
||||
For embedding retrieval, you can employ the BGE-M3 model using the same approach as BGE.
|
||||
The only difference is that the BGE-M3 model no longer requires adding instructions to the queries.
|
||||
|
||||
For hybrid retrieval, you can use [Vespa](https://github.com/vespa-engine/pyvespa/blob/master/docs/sphinx/source/examples/mother-of-all-embedding-models-cloud.ipynb
|
||||
) and [Milvus](https://github.com/milvus-io/pymilvus/blob/master/examples/hello_hybrid_sparse_dense.py).
|
||||
|
||||
|
||||
**3. How to fine-tune bge-M3 model?**
|
||||
|
||||
You can follow the common in this [example](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/finetune/embedder#1-standard-model)
|
||||
to fine-tune the dense embedding.
|
||||
|
||||
If you want to fine-tune all embedding function of m3 (dense, sparse and colbert), you can refer to the [unified_fine-tuning example](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/finetune/embedder#2-bge-m3)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
## Usage
|
||||
|
||||
Install:
|
||||
```
|
||||
git clone https://github.com/FlagOpen/FlagEmbedding.git
|
||||
cd FlagEmbedding
|
||||
pip install -e .
|
||||
```
|
||||
or:
|
||||
```
|
||||
pip install -U FlagEmbedding
|
||||
```
|
||||
|
||||
|
||||
|
||||
### Generate Embedding for text
|
||||
|
||||
- Dense Embedding
|
||||
```python
|
||||
from FlagEmbedding import BGEM3FlagModel
|
||||
|
||||
model = BGEM3FlagModel('BAAI/bge-m3',
|
||||
use_fp16=True,
|
||||
devices=['cuda:0']) # Setting use_fp16 to True speeds up computation with a slight performance degradation
|
||||
|
||||
sentences_1 = ["What is BGE M3?", "Defination of BM25"]
|
||||
sentences_2 = ["BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
|
||||
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"]
|
||||
|
||||
embeddings_1 = model.encode(sentences_1,
|
||||
batch_size=12,
|
||||
max_length=8192, # If you don't need such a long length, you can set a smaller value to speed up the encoding process.
|
||||
)['dense_vecs']
|
||||
embeddings_2 = model.encode(sentences_2)['dense_vecs']
|
||||
similarity = embeddings_1 @ embeddings_2.T
|
||||
print(similarity)
|
||||
# [[0.6265, 0.3477], [0.3499, 0.678 ]]
|
||||
```
|
||||
You also can use sentence-transformers and huggingface transformers to generate dense embeddings.
|
||||
Refer to [baai_general_embedding](https://github.com/FlagOpen/FlagEmbedding/tree/master/FlagEmbedding/baai_general_embedding#usage) for details.
|
||||
|
||||
|
||||
- Sparse Embedding (Lexical Weight)
|
||||
```python
|
||||
from FlagEmbedding import BGEM3FlagModel
|
||||
|
||||
model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation
|
||||
|
||||
sentences_1 = ["What is BGE M3?", "Defination of BM25"]
|
||||
sentences_2 = ["BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
|
||||
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"]
|
||||
|
||||
output_1 = model.encode(sentences_1, return_dense=True, return_sparse=True, return_colbert_vecs=False)
|
||||
output_2 = model.encode(sentences_2, return_dense=True, return_sparse=True, return_colbert_vecs=False)
|
||||
|
||||
# you can see the weight for each token:
|
||||
print(model.convert_id_to_token(output_1['lexical_weights']))
|
||||
# [{'What': 0.08356, 'is': 0.0814, 'B': 0.1296, 'GE': 0.252, 'M': 0.1702, '3': 0.2695, '?': 0.04092},
|
||||
# {'De': 0.05005, 'fin': 0.1368, 'ation': 0.04498, 'of': 0.0633, 'BM': 0.2515, '25': 0.3335}]
|
||||
|
||||
|
||||
# compute the scores via lexical mathcing
|
||||
lexical_scores = model.compute_lexical_matching_score(output_1['lexical_weights'][0], output_2['lexical_weights'][0])
|
||||
print(lexical_scores)
|
||||
# 0.19554901123046875
|
||||
|
||||
print(model.compute_lexical_matching_score(output_1['lexical_weights'][0], output_1['lexical_weights'][1]))
|
||||
# 0.0
|
||||
```
|
||||
|
||||
- Multi-Vector (ColBERT)
|
||||
```python
|
||||
from FlagEmbedding import BGEM3FlagModel
|
||||
|
||||
model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True)
|
||||
|
||||
sentences_1 = ["What is BGE M3?", "Defination of BM25"]
|
||||
sentences_2 = ["BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
|
||||
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"]
|
||||
|
||||
output_1 = model.encode(sentences_1, return_dense=True, return_sparse=True, return_colbert_vecs=True)
|
||||
output_2 = model.encode(sentences_2, return_dense=True, return_sparse=True, return_colbert_vecs=True)
|
||||
|
||||
print(model.colbert_score(output_1['colbert_vecs'][0], output_2['colbert_vecs'][0]))
|
||||
print(model.colbert_score(output_1['colbert_vecs'][0], output_2['colbert_vecs'][1]))
|
||||
# 0.7797
|
||||
# 0.4620
|
||||
```
|
||||
|
||||
|
||||
### Compute score for text pairs
|
||||
Input a list of text pairs, you can get the scores computed by different methods.
|
||||
```python
|
||||
from FlagEmbedding import BGEM3FlagModel
|
||||
|
||||
model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True)
|
||||
|
||||
sentences_1 = ["What is BGE M3?", "Defination of BM25"]
|
||||
sentences_2 = ["BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
|
||||
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"]
|
||||
|
||||
sentence_pairs = [[i,j] for i in sentences_1 for j in sentences_2]
|
||||
|
||||
print(model.compute_score(sentence_pairs,
|
||||
max_passage_length=128, # a smaller max length leads to a lower latency
|
||||
weights_for_different_modes=[0.4, 0.2, 0.4])) # weights_for_different_modes(w) is used to do weighted sum: w[0]*dense_score + w[1]*sparse_score + w[2]*colbert_score
|
||||
|
||||
# {
|
||||
# 'colbert': [0.7796499729156494, 0.4621465802192688, 0.4523794651031494, 0.7898575067520142],
|
||||
# 'sparse': [0.195556640625, 0.00879669189453125, 0.0, 0.1802978515625],
|
||||
# 'dense': [0.6259765625, 0.347412109375, 0.349853515625, 0.67822265625],
|
||||
# 'sparse+dense': [0.482503205537796, 0.23454029858112335, 0.2332356721162796, 0.5122477412223816],
|
||||
# 'colbert+sparse+dense': [0.6013619303703308, 0.3255828022956848, 0.32089319825172424, 0.6232916116714478]
|
||||
# }
|
||||
```
|
||||
|
||||
|
||||
|
||||
|
||||
## Evaluation
|
||||
|
||||
We provide the evaluation script for [MKQA](https://github.com/FlagOpen/FlagEmbedding/tree/master/C_MTEB/MKQA) and [MLDR](https://github.com/FlagOpen/FlagEmbedding/tree/master/C_MTEB/MLDR)
|
||||
|
||||
|
||||
### Benchmarks from the open-source community
|
||||
|
||||
The BGE-M3 model emerged as the top performer on this benchmark (OAI is short for OpenAI).
|
||||
For more details, please refer to the [article](https://towardsdatascience.com/openai-vs-open-source-multilingual-embedding-models-e5ccb7c90f05) and [Github Repo](https://github.com/Yannael/multilingual-embeddings)
|
||||
|
||||
|
||||
### Our results
|
||||
- Multilingual (MIRACL dataset)
|
||||
|
||||
|
||||
|
||||
- Cross-lingual (MKQA dataset)
|
||||
|
||||
|
||||
|
||||
- Long Document Retrieval
|
||||
- MLDR:
|
||||
|
||||
Please note that [MLDR](https://huggingface.co/datasets/Shitao/MLDR) is a document retrieval dataset we constructed via LLM,
|
||||
covering 13 languages, including test set, validation set, and training set.
|
||||
We utilized the training set from MLDR to enhance the model's long document retrieval capabilities.
|
||||
Therefore, comparing baselines with `Dense w.o.long`(fine-tuning without long document dataset) is more equitable.
|
||||
Additionally, this long document retrieval dataset will be open-sourced to address the current lack of open-source multilingual long text retrieval datasets.
|
||||
We believe that this data will be helpful for the open-source community in training document retrieval models.
|
||||
|
||||
- NarritiveQA:
|
||||
|
||||
|
||||
- Comparison with BM25
|
||||
|
||||
We utilized Pyserini to implement BM25, and the test results can be reproduced by this [script](https://github.com/FlagOpen/FlagEmbedding/tree/master/C_MTEB/MLDR#bm25-baseline).
|
||||
We tested BM25 using two different tokenizers:
|
||||
one using Lucene Analyzer and the other using the same tokenizer as M3 (i.e., the tokenizer of xlm-roberta).
|
||||
The results indicate that BM25 remains a competitive baseline,
|
||||
especially in long document retrieval.
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
## Training
|
||||
- Self-knowledge Distillation: combining multiple outputs from different
|
||||
retrieval modes as reward signal to enhance the performance of single mode(especially for sparse retrieval and multi-vec(colbert) retrival)
|
||||
- Efficient Batching: Improve the efficiency when fine-tuning on long text.
|
||||
The small-batch strategy is simple but effective, which also can used to fine-tune large embedding model.
|
||||
- MCLS: A simple method to improve the performance on long text without fine-tuning.
|
||||
If you have no enough resource to fine-tuning model with long text, the method is useful.
|
||||
|
||||
Refer to our [report](https://arxiv.org/pdf/2402.03216.pdf) for more details.
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
## Acknowledgement
|
||||
|
||||
Thanks to the authors of open-sourced datasets, including Miracl, MKQA, NarritiveQA, etc.
|
||||
Thanks to the open-sourced libraries like [Tevatron](https://github.com/texttron/tevatron), [Pyserini](https://github.com/castorini/pyserini).
|
||||
|
||||
|
||||
|
||||
## Citation
|
||||
|
||||
If you find this repository useful, please consider giving a star :star: and citation
|
||||
|
||||
```
|
||||
@misc{bge-m3,
|
||||
title={BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation},
|
||||
author={Jianlv Chen and Shitao Xiao and Peitian Zhang and Kun Luo and Defu Lian and Zheng Liu},
|
||||
year={2024},
|
||||
eprint={2402.03216},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CL}
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
from .modeling import BGEM3Model, BGEM3ForInference, EncoderOutput
|
||||
from .trainer import BiTrainer
|
||||
@@ -0,0 +1,100 @@
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
from transformers import TrainingArguments
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArguments:
|
||||
"""
|
||||
Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
|
||||
"""
|
||||
|
||||
model_name_or_path: str = field(
|
||||
metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
|
||||
)
|
||||
config_name: Optional[str] = field(
|
||||
default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
|
||||
)
|
||||
tokenizer_name: Optional[str] = field(
|
||||
default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
|
||||
)
|
||||
cache_dir: Optional[str] = field(
|
||||
default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataArguments:
|
||||
knowledge_distillation: bool = field(
|
||||
default=False, metadata={"help": "Use knowledge distillation when `pos_scores` and `neg_scores` are in features of training data"}
|
||||
)
|
||||
train_data: str = field(
|
||||
default=None, metadata={"help": "One or more paths to training data", "nargs": "+"}
|
||||
)
|
||||
cache_path: Optional[str] = field(
|
||||
default=None, metadata={"help": "Where do you want to store the cached data"}
|
||||
)
|
||||
train_group_size: int = field(default=8)
|
||||
|
||||
query_max_len: int = field(
|
||||
default=32,
|
||||
metadata={
|
||||
"help": "The maximum total input sequence length after tokenization for passage. Sequences longer "
|
||||
"than this will be truncated, sequences shorter will be padded."
|
||||
},
|
||||
)
|
||||
|
||||
passage_max_len: int = field(
|
||||
default=128,
|
||||
metadata={
|
||||
"help": "The maximum total input sequence length after tokenization for passage. Sequences longer "
|
||||
"than this will be truncated, sequences shorter will be padded."
|
||||
},
|
||||
)
|
||||
|
||||
max_example_num_per_dataset: int = field(
|
||||
default=None, metadata={"help": "the max number of examples for each dataset"}
|
||||
)
|
||||
|
||||
query_instruction_for_retrieval: str= field(
|
||||
default=None, metadata={"help": "instruction for query"}
|
||||
)
|
||||
passage_instruction_for_retrieval: str = field(
|
||||
default=None, metadata={"help": "instruction for passage"}
|
||||
)
|
||||
|
||||
same_task_within_batch: bool = field(
|
||||
default=False, metadata={"help": "All samples in the same batch comes from the same task."}
|
||||
)
|
||||
shuffle_ratio: float = field(
|
||||
default=0.0, metadata={"help": "The ratio of shuffling the text"}
|
||||
)
|
||||
|
||||
small_threshold: int = field(
|
||||
default=0, metadata={"help": "The threshold of small dataset. All small dataset in the same directory will be merged into one dataset."}
|
||||
)
|
||||
drop_threshold: int = field(
|
||||
default=0, metadata={"help": "The threshold for dropping merged small dataset. If the number of examples in the merged small dataset is less than this threshold, it will be dropped."}
|
||||
)
|
||||
|
||||
def __post_init__(self):
|
||||
for train_dir in self.train_data:
|
||||
if not os.path.exists(train_dir):
|
||||
raise FileNotFoundError(f"cannot find file: {train_dir}, please set a true path")
|
||||
|
||||
@dataclass
|
||||
class RetrieverTrainingArguments(TrainingArguments):
|
||||
negatives_cross_device: bool = field(default=False, metadata={"help": "share negatives across devices"})
|
||||
temperature: Optional[float] = field(default=0.02)
|
||||
fix_position_embedding: bool = field(default=False, metadata={"help": "Freeze the parameters of position embeddings"})
|
||||
sentence_pooling_method: str = field(default='cls', metadata={"help": "the pooling method, should be cls or mean"})
|
||||
normlized: bool = field(default=True)
|
||||
enable_sub_batch: bool = field(default=True, metadata={"help": "Freeze the parameters of position embeddings"})
|
||||
|
||||
unified_finetuning: bool = field(default=False, metadata={"help": "use unify fine-tuning"})
|
||||
use_self_distill: bool = field(default=False, metadata={"help": "use self-distill when using unify fine-tuning"})
|
||||
fix_encoder: bool = field(default=False, metadata={"help": "Freeze the parameters of encoder"})
|
||||
colbert_dim: int = field(default=-1, metadata={"help": "Dim of colbert linear"})
|
||||
self_distill_start_step: int = field(default=-1, metadata={"help": "Num of step when using self-distill"})
|
||||
@@ -0,0 +1,301 @@
|
||||
import math
|
||||
import os.path
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
import torch
|
||||
import numpy as np
|
||||
import datasets
|
||||
from pprint import pprint
|
||||
from torch.utils.data import Dataset
|
||||
from transformers import DataCollatorWithPadding
|
||||
import torch.distributed as dist
|
||||
|
||||
from .arguments import DataArguments
|
||||
|
||||
|
||||
class SameDatasetTrainDataset(Dataset):
|
||||
"""Dataset to yield a batch of data at one time. All samples in the same batch comes from the same task.
|
||||
"""
|
||||
def __init__(self, args: DataArguments, batch_size: int, seed: int, process_index: int=0, num_processes: int=1):
|
||||
train_datasets = []
|
||||
each_data_inxs = []
|
||||
batch_size_inxs = []
|
||||
pqloss_flag = []
|
||||
cur_all_num = 0
|
||||
|
||||
SMALL_THRESHOLD = args.small_threshold
|
||||
DROP_THRESHOLD = args.drop_threshold
|
||||
|
||||
context_feat = datasets.Features({
|
||||
'query': datasets.Value('string'),
|
||||
'pos': datasets.Sequence(datasets.Value('string')),
|
||||
'neg': datasets.Sequence(datasets.Value('string'))
|
||||
})
|
||||
context_feat_kd = datasets.Features({
|
||||
'query': datasets.Value('string'),
|
||||
'pos': datasets.Sequence(datasets.Value('string')),
|
||||
'neg': datasets.Sequence(datasets.Value('string')),
|
||||
'pos_scores': datasets.Sequence(datasets.Value('float')),
|
||||
'neg_scores': datasets.Sequence(datasets.Value('float')),
|
||||
})
|
||||
assert isinstance(args.train_data, list) and len(args.train_data) >= 1
|
||||
|
||||
if dist.get_rank() == 0:
|
||||
self.print_batch_size(batch_size=batch_size, train_group_size=args.train_group_size)
|
||||
|
||||
for data_dir in args.train_data:
|
||||
if not os.path.isdir(data_dir):
|
||||
raise FileNotFoundError(f"{data_dir} is a file, not a directionary")
|
||||
|
||||
small_datasets = []
|
||||
small_batch_size = math.inf
|
||||
|
||||
# Add `parallel_` in `data_dir` to indicate that this dataset is parallel corpus
|
||||
flag = 'parallel_' in data_dir
|
||||
for file in os.listdir(data_dir):
|
||||
if not (file.endswith('.json') or file.endswith('.jsonl')):
|
||||
continue
|
||||
|
||||
file_path = os.path.join(data_dir, file)
|
||||
if dist.get_rank() == 0:
|
||||
print(f'loading data from {file_path} ...')
|
||||
try:
|
||||
temp_dataset = datasets.load_dataset('json', data_files=file_path, split='train', cache_dir=args.cache_path, features=context_feat)
|
||||
except:
|
||||
temp_dataset = datasets.load_dataset('json', data_files=file_path, split='train', cache_dir=args.cache_path, features=context_feat_kd)
|
||||
if not args.knowledge_distillation:
|
||||
temp_dataset = temp_dataset.remove_columns(['pos_scores', 'neg_scores'])
|
||||
|
||||
if len(temp_dataset) == 0:
|
||||
continue
|
||||
elif len(temp_dataset) < SMALL_THRESHOLD:
|
||||
small_datasets.append(temp_dataset)
|
||||
small_batch_size = min(small_batch_size, self.get_file_batch_size(file, batch_size, train_group_size=args.train_group_size))
|
||||
else:
|
||||
if args.max_example_num_per_dataset is not None and len(temp_dataset) > args.max_example_num_per_dataset:
|
||||
temp_dataset = temp_dataset.select(
|
||||
random.sample(list(range(len(temp_dataset))), args.max_example_num_per_dataset))
|
||||
train_datasets.append(temp_dataset)
|
||||
each_data_inxs.append(np.arange(len(temp_dataset)) + cur_all_num)
|
||||
cur_all_num += len(temp_dataset)
|
||||
batch_size_inxs.append(self.get_file_batch_size(file, batch_size, train_group_size=args.train_group_size))
|
||||
pqloss_flag.append(flag)
|
||||
|
||||
if len(small_datasets) > 0:
|
||||
small_dataset = datasets.concatenate_datasets(small_datasets)
|
||||
if len(small_dataset) >= DROP_THRESHOLD:
|
||||
train_datasets.append(small_dataset)
|
||||
each_data_inxs.append(np.arange(len(small_dataset)) + cur_all_num)
|
||||
cur_all_num += len(small_dataset)
|
||||
batch_size_inxs.append(small_batch_size)
|
||||
pqloss_flag.append(flag)
|
||||
|
||||
self.dataset = datasets.concatenate_datasets(train_datasets)
|
||||
self.each_data_inxs = each_data_inxs
|
||||
self.datasets_inxs = np.arange(len(each_data_inxs))
|
||||
self.batch_size_inxs = batch_size_inxs
|
||||
self.pqloss_flag = pqloss_flag
|
||||
|
||||
self.process_index = process_index
|
||||
self.num_processes = num_processes
|
||||
self.args = args
|
||||
self.shuffle_ratio = args.shuffle_ratio
|
||||
|
||||
self.deterministic_generator = np.random.default_rng(seed)
|
||||
self.step = 0
|
||||
self.refresh_epoch()
|
||||
|
||||
def print_batch_size(self, batch_size: int, train_group_size: int):
|
||||
length_list = ['0-500', '500-1000', '1000-2000', '2000-3000', '3000-4000', '4000-5000', '5000-6000', '6000-7000', '7000-inf']
|
||||
batch_size_dict = {
|
||||
k: self.get_file_batch_size(f"len-{k}.jsonl", batch_size, train_group_size) for k in length_list
|
||||
}
|
||||
batch_size_list = [
|
||||
f'{length}: {batch_size_dict[length]}' for length in length_list
|
||||
]
|
||||
print("=========================")
|
||||
print("Batch Size Dict:")
|
||||
pprint(batch_size_list)
|
||||
print("=========================")
|
||||
|
||||
@staticmethod
|
||||
def get_file_batch_size(file: str, batch_size: int, train_group_size: int):
|
||||
if train_group_size == 8:
|
||||
# 80GB
|
||||
if 'len-0-500.jsonl' in file:
|
||||
return 48
|
||||
elif 'len-500-1000.jsonl' in file:
|
||||
return 32
|
||||
elif 'len-1000-2000.jsonl' in file:
|
||||
return 20
|
||||
elif 'len-2000-3000.jsonl' in file:
|
||||
return 18
|
||||
elif 'len-3000-4000.jsonl' in file:
|
||||
return 14
|
||||
elif 'len-4000-5000.jsonl' in file:
|
||||
return 14
|
||||
elif 'len-5000-6000.jsonl' in file:
|
||||
return 12
|
||||
elif 'len-6000-7000.jsonl' in file:
|
||||
return 10
|
||||
elif 'len-7000-inf.jsonl' in file:
|
||||
return 8
|
||||
else:
|
||||
return batch_size
|
||||
elif train_group_size == 1:
|
||||
# 80GB
|
||||
if 'len-0-500.jsonl' in file:
|
||||
return 700
|
||||
elif 'len-500-1000.jsonl' in file:
|
||||
return 570
|
||||
elif 'len-1000-2000.jsonl' in file:
|
||||
return 388
|
||||
elif 'len-2000-3000.jsonl' in file:
|
||||
return 288
|
||||
elif 'len-3000-4000.jsonl' in file:
|
||||
return 224
|
||||
elif 'len-4000-5000.jsonl' in file:
|
||||
return 180
|
||||
elif 'len-5000-6000.jsonl' in file:
|
||||
return 157
|
||||
elif 'len-6000-7000.jsonl' in file:
|
||||
return 128
|
||||
elif 'len-7000-inf.jsonl' in file:
|
||||
return 104
|
||||
else:
|
||||
return batch_size
|
||||
else:
|
||||
return batch_size
|
||||
|
||||
def refresh_epoch(self):
|
||||
print(f'---------------------------*Rank {self.process_index}: refresh data---------------------------')
|
||||
self.deterministic_generator.shuffle(self.datasets_inxs)
|
||||
# Dynamically adjust batch size
|
||||
batch_datas = []
|
||||
for dataset_inx in self.datasets_inxs:
|
||||
self.deterministic_generator.shuffle(self.each_data_inxs[dataset_inx])
|
||||
cur_batch_size = self.batch_size_inxs[dataset_inx]*self.num_processes
|
||||
flag = self.pqloss_flag[dataset_inx]
|
||||
for start_index in range(0, len(self.each_data_inxs[dataset_inx]), cur_batch_size):
|
||||
# judge the last batch's length
|
||||
if len(self.each_data_inxs[dataset_inx]) - start_index < 2 * self.num_processes:
|
||||
break
|
||||
batch_datas.append((self.each_data_inxs[dataset_inx][start_index:start_index+cur_batch_size], flag))
|
||||
self.deterministic_generator.shuffle(batch_datas)
|
||||
self.batch_datas = batch_datas
|
||||
self.step = 0
|
||||
|
||||
def __getitem__(self, _):
|
||||
batch_indices, pqloss_flag = self.batch_datas[self.step]
|
||||
cur_batch_size = int(len(batch_indices) / self.num_processes)
|
||||
batch_indices = batch_indices[self.process_index * cur_batch_size: (self.process_index + 1) * cur_batch_size]
|
||||
batch_data = self.dataset[batch_indices]
|
||||
self.step += 1
|
||||
queries, passages, teacher_scores = self.create_batch_data(batch_raw_data=batch_data)
|
||||
# print('rank, step, flag, query, passage:', dist.get_rank(), self.step, pqloss_flag, queries, passages)
|
||||
return queries, passages, teacher_scores, pqloss_flag
|
||||
|
||||
def shuffle_text(self, text):
|
||||
if self.shuffle_ratio > 0 and len(text) > 100 and random.random() < self.shuffle_ratio:
|
||||
split_text = []
|
||||
chunk_size = len(text)//3 + 1
|
||||
for i in range(0, len(text), chunk_size):
|
||||
split_text.append(text[i:i+chunk_size])
|
||||
random.shuffle(split_text)
|
||||
return " ".join(split_text)
|
||||
else:
|
||||
return text
|
||||
|
||||
def create_batch_data(self, batch_raw_data):
|
||||
queries, passages = [], []
|
||||
teacher_scores = []
|
||||
for i in range(len(batch_raw_data['query'])):
|
||||
queries.append(batch_raw_data['query'][i])
|
||||
|
||||
pos_inx = random.choice(list(range(len(batch_raw_data['pos'][i]))))
|
||||
passages.append(self.shuffle_text(batch_raw_data['pos'][i][pos_inx]))
|
||||
if 'pos_scores' in batch_raw_data and batch_raw_data['pos_scores'][i] is not None:
|
||||
teacher_scores.append(batch_raw_data['pos_scores'][i][pos_inx])
|
||||
|
||||
neg_inx_set = list(range(len(batch_raw_data['neg'][i])))
|
||||
if len(batch_raw_data['neg'][i]) < self.args.train_group_size - 1:
|
||||
num = math.ceil((self.args.train_group_size - 1) / len(batch_raw_data['neg'][i]))
|
||||
neg_inxs = random.sample(neg_inx_set * num, self.args.train_group_size - 1)
|
||||
else:
|
||||
neg_inxs = random.sample(neg_inx_set, self.args.train_group_size - 1)
|
||||
|
||||
if 'neg_scores' in batch_raw_data and batch_raw_data['neg_scores'][i] is not None:
|
||||
neg_scores = [(x, batch_raw_data['neg_scores'][i][x]) for x in neg_inxs]
|
||||
neg_scores = sorted(neg_scores, key=lambda x:x[1], reverse=True)
|
||||
neg_inxs = [x[0] for x in neg_scores]
|
||||
teacher_scores.extend([x[1] for x in neg_scores])
|
||||
|
||||
negs = [batch_raw_data['neg'][i][x] for x in neg_inxs]
|
||||
passages.extend(negs)
|
||||
|
||||
if len(teacher_scores) > 0 and len(passages) > 0:
|
||||
assert len(teacher_scores) == len(passages)
|
||||
|
||||
if self.args.query_instruction_for_retrieval is not None:
|
||||
queries = [self.args.query_instruction_for_retrieval+q for q in queries]
|
||||
if self.args.passage_instruction_for_retrieval is not None:
|
||||
passages = [self.args.passage_instruction_for_retrieval+p for p in passages]
|
||||
|
||||
if len(teacher_scores) == 0:
|
||||
teacher_scores = None
|
||||
return queries, passages, teacher_scores
|
||||
|
||||
def __len__(self):
|
||||
return len(self.batch_datas) * self.num_processes
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmbedCollator(DataCollatorWithPadding):
|
||||
"""
|
||||
Wrapper that does conversion from List[Tuple[encode_qry, encode_psg]] to List[qry], List[psg]
|
||||
and pass batch separately to the actual collator.
|
||||
Abstract out data detail for the model.
|
||||
"""
|
||||
query_max_len: int = 32
|
||||
passage_max_len: int = 128
|
||||
|
||||
def __call__(self, features):
|
||||
query = [f[0] for f in features]
|
||||
passage = [f[1] for f in features]
|
||||
|
||||
teacher_scores = None
|
||||
if len(features[0]) > 2:
|
||||
teacher_scores = [f[2] for f in features]
|
||||
if teacher_scores[0] is None:
|
||||
teacher_scores = None
|
||||
else:
|
||||
teacher_scores = torch.FloatTensor(teacher_scores)
|
||||
|
||||
flag = None
|
||||
if len(features[0]) == 4:
|
||||
flag = [f[3] for f in features][0]
|
||||
|
||||
if isinstance(query[0], list):
|
||||
query = sum(query, [])
|
||||
if isinstance(passage[0], list):
|
||||
passage = sum(passage, [])
|
||||
|
||||
q_collated = self.tokenizer(
|
||||
query,
|
||||
# padding='max_length', # used for adjusting the batch size in `get_file_batch_size()`
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=self.query_max_len,
|
||||
return_tensors="pt",
|
||||
)
|
||||
d_collated = self.tokenizer(
|
||||
passage,
|
||||
# padding='max_length', # used for adjusting the batch size in `get_file_batch_size()`
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=self.passage_max_len,
|
||||
return_tensors="pt",
|
||||
)
|
||||
if teacher_scores is not None:
|
||||
teacher_scores = teacher_scores.reshape((len(q_collated['input_ids']), -1))
|
||||
return {"query": q_collated, "passage": d_collated, "teacher_scores": teacher_scores, "bi_directions": flag}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 129 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 474 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 563 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 594 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 155 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 20 KiB |
@@ -0,0 +1,390 @@
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, Optional
|
||||
import os
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch import nn, Tensor
|
||||
import torch.nn.functional as F
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
from transformers.file_utils import ModelOutput
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EncoderOutput(ModelOutput):
|
||||
q_reps: Optional[Tensor] = None
|
||||
p_reps: Optional[Tensor] = None
|
||||
loss: Optional[Tensor] = None
|
||||
scores: Optional[Tensor] = None
|
||||
|
||||
|
||||
class BGEM3Model(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
model_name: str = None,
|
||||
normlized: bool = True,
|
||||
sentence_pooling_method: str = 'cls',
|
||||
negatives_cross_device: bool = False,
|
||||
temperature: float = 1.0,
|
||||
enable_sub_batch: bool = True,
|
||||
unified_finetuning: bool = True,
|
||||
use_self_distill: bool = False,
|
||||
colbert_dim: int = -1,
|
||||
self_distill_start_step: int = -1,
|
||||
):
|
||||
super().__init__()
|
||||
self.load_model(model_name, colbert_dim=colbert_dim)
|
||||
self.vocab_size = self.model.config.vocab_size
|
||||
self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')
|
||||
|
||||
self.unified_finetuning = unified_finetuning
|
||||
if not self.unified_finetuning:
|
||||
self.colbert_linear = None
|
||||
self.sparse_linear = None
|
||||
|
||||
self.normlized = normlized
|
||||
self.sentence_pooling_method = sentence_pooling_method
|
||||
self.enable_sub_batch = enable_sub_batch
|
||||
self.temperature = temperature
|
||||
self.use_self_distill = use_self_distill
|
||||
self.self_distill_start_step = self_distill_start_step
|
||||
|
||||
self.step = 0
|
||||
if not normlized:
|
||||
self.temperature = 1.0
|
||||
logger.info("reset temperature = 1.0 due to using inner product to compute similarity")
|
||||
|
||||
self.negatives_cross_device = negatives_cross_device
|
||||
if self.negatives_cross_device:
|
||||
if not dist.is_initialized():
|
||||
raise ValueError('Distributed training has not been initialized for representation all gather.')
|
||||
|
||||
self.process_rank = dist.get_rank()
|
||||
self.world_size = dist.get_world_size()
|
||||
|
||||
def load_model(self, model_name, colbert_dim: int = -1):
|
||||
if not os.path.exists(model_name):
|
||||
cache_folder = os.getenv('HF_HUB_CACHE')
|
||||
model_name = snapshot_download(repo_id=model_name,
|
||||
cache_dir=cache_folder,
|
||||
ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5'])
|
||||
|
||||
self.model = AutoModel.from_pretrained(model_name)
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
|
||||
self.colbert_linear = torch.nn.Linear(in_features=self.model.config.hidden_size,
|
||||
out_features=self.model.config.hidden_size if colbert_dim == -1 else colbert_dim)
|
||||
self.sparse_linear = torch.nn.Linear(in_features=self.model.config.hidden_size, out_features=1)
|
||||
|
||||
if os.path.exists(os.path.join(model_name, 'colbert_linear.pt')) and os.path.exists(
|
||||
os.path.join(model_name, 'sparse_linear.pt')):
|
||||
logger.info('loading existing colbert_linear and sparse_linear---------')
|
||||
self.load_pooler(model_dir=model_name)
|
||||
else:
|
||||
logger.info(
|
||||
'The parameters of colbert_linear and sparse linear is new initialize. Make sure the model is loaded for training, not inferencing')
|
||||
|
||||
def gradient_checkpointing_enable(self, **kwargs):
|
||||
self.model.gradient_checkpointing_enable(**kwargs)
|
||||
|
||||
def dense_embedding(self, hidden_state, mask):
|
||||
if self.sentence_pooling_method == 'cls':
|
||||
return hidden_state[:, 0]
|
||||
elif self.sentence_pooling_method == 'mean':
|
||||
s = torch.sum(hidden_state * mask.unsqueeze(-1).float(), dim=1)
|
||||
d = mask.sum(axis=1, keepdim=True).float()
|
||||
return s / d
|
||||
|
||||
def sparse_embedding(self, hidden_state, input_ids, return_embedding: bool = True):
|
||||
token_weights = torch.relu(self.sparse_linear(hidden_state))
|
||||
if not return_embedding: return token_weights
|
||||
|
||||
if self.training:
|
||||
sparse_embedding = torch.zeros(
|
||||
input_ids.size(0), input_ids.size(1), self.vocab_size,
|
||||
dtype=token_weights.dtype,
|
||||
device=token_weights.device
|
||||
)
|
||||
sparse_embedding = torch.scatter(sparse_embedding, dim=-1, index=input_ids.unsqueeze(-1), src=token_weights)
|
||||
sparse_embedding = torch.max(sparse_embedding, dim=1).values
|
||||
else:
|
||||
# Optimize suggestion from issue #1364: https://github.com/FlagOpen/FlagEmbedding/issues/1364
|
||||
# Disable when self.training = True, otherwise will cause:
|
||||
# RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation
|
||||
sparse_embedding = torch.zeros(
|
||||
input_ids.size(0), self.vocab_size,
|
||||
dtype=token_weights.dtype,
|
||||
device=token_weights.device
|
||||
)
|
||||
sparse_embedding = sparse_embedding.scatter_reduce(
|
||||
dim=-1, index=input_ids, src=token_weights.squeeze(-1), reduce="amax"
|
||||
)
|
||||
|
||||
unused_tokens = [self.tokenizer.cls_token_id, self.tokenizer.eos_token_id, self.tokenizer.pad_token_id,
|
||||
self.tokenizer.unk_token_id]
|
||||
sparse_embedding[:, unused_tokens] *= 0.
|
||||
return sparse_embedding
|
||||
|
||||
def colbert_embedding(self, last_hidden_state, mask):
|
||||
colbert_vecs = self.colbert_linear(last_hidden_state[:, 1:])
|
||||
colbert_vecs = colbert_vecs * mask[:, 1:][:, :, None].float()
|
||||
return colbert_vecs
|
||||
|
||||
def dense_score(self, q_reps, p_reps):
|
||||
scores = self.compute_similarity(q_reps, p_reps) / self.temperature
|
||||
scores = scores.view(q_reps.size(0), -1)
|
||||
return scores
|
||||
|
||||
def sparse_score(self, q_reps, p_reps):
|
||||
scores = self.compute_similarity(q_reps, p_reps) / self.temperature
|
||||
scores = scores.view(q_reps.size(0), -1)
|
||||
return scores
|
||||
|
||||
def colbert_score(self, q_reps, p_reps, q_mask: torch.Tensor):
|
||||
token_scores = torch.einsum('qin,pjn->qipj', q_reps, p_reps)
|
||||
scores, _ = token_scores.max(-1)
|
||||
scores = scores.sum(1) / q_mask[:, 1:].sum(-1, keepdim=True)
|
||||
scores = scores / self.temperature
|
||||
return scores
|
||||
|
||||
def _encode(self, features):
|
||||
dense_vecs, sparse_vecs, colbert_vecs = None, None, None
|
||||
last_hidden_state = self.model(**features, return_dict=True).last_hidden_state
|
||||
dense_vecs = self.dense_embedding(last_hidden_state, features['attention_mask'])
|
||||
if self.unified_finetuning:
|
||||
sparse_vecs = self.sparse_embedding(last_hidden_state, features['input_ids'])
|
||||
colbert_vecs = self.colbert_embedding(last_hidden_state, features['attention_mask'])
|
||||
if self.normlized:
|
||||
dense_vecs = torch.nn.functional.normalize(dense_vecs, dim=-1)
|
||||
if self.unified_finetuning:
|
||||
colbert_vecs = torch.nn.functional.normalize(colbert_vecs, dim=-1)
|
||||
return dense_vecs, sparse_vecs, colbert_vecs
|
||||
|
||||
def encode(self, features, sub_batch_size=None):
|
||||
if features is None:
|
||||
return None
|
||||
|
||||
if sub_batch_size is not None and sub_batch_size != -1:
|
||||
all_dense_vecs, all_sparse_vecs, all_colbert_vecs = [], [], []
|
||||
for i in range(0, len(features['attention_mask']), sub_batch_size):
|
||||
end_inx = min(i + sub_batch_size, len(features['attention_mask']))
|
||||
sub_features = {}
|
||||
for k, v in features.items():
|
||||
sub_features[k] = v[i:end_inx]
|
||||
|
||||
dense_vecs, sparse_vecs, colbert_vecs = self._encode(sub_features)
|
||||
all_dense_vecs.append(dense_vecs)
|
||||
all_sparse_vecs.append(sparse_vecs)
|
||||
all_colbert_vecs.append(colbert_vecs)
|
||||
|
||||
dense_vecs = torch.cat(all_dense_vecs, 0)
|
||||
if self.unified_finetuning:
|
||||
sparse_vecs = torch.cat(all_sparse_vecs, 0)
|
||||
colbert_vecs = torch.cat(all_colbert_vecs, 0)
|
||||
else:
|
||||
dense_vecs, sparse_vecs, colbert_vecs = self._encode(features)
|
||||
|
||||
if self.unified_finetuning:
|
||||
return dense_vecs.contiguous(), sparse_vecs.contiguous(), colbert_vecs.contiguous()
|
||||
else:
|
||||
return dense_vecs.contiguous(), None, None
|
||||
|
||||
def compute_sub_batch_size(self, features):
|
||||
mapping = [(6000, 1), (5000, 2), (4000, 3), (3000, 3), (2000, 5), (1000, 9), (512, 16), (0, 32)]
|
||||
cur_l = features['input_ids'].size(-1)
|
||||
for l, b in mapping:
|
||||
if cur_l >= l:
|
||||
return b
|
||||
|
||||
def compute_similarity(self, q_reps, p_reps):
|
||||
if len(p_reps.size()) == 2:
|
||||
return torch.matmul(q_reps, p_reps.transpose(0, 1))
|
||||
return torch.matmul(q_reps, p_reps.transpose(-2, -1))
|
||||
|
||||
def distill_loss(self, teacher_targets, student_scores, group_size):
|
||||
labels = torch.arange(student_scores.size(0), device=student_scores.device, dtype=torch.long)
|
||||
labels = labels * group_size
|
||||
|
||||
loss = 0
|
||||
mask = torch.zeros_like(student_scores)
|
||||
for i in range(group_size):
|
||||
temp_target = labels + i
|
||||
temp_scores = student_scores + mask
|
||||
temp_loss = F.cross_entropy(temp_scores, temp_target, reduction="none") # B
|
||||
loss += torch.mean(teacher_targets[:, i] * temp_loss)
|
||||
mask = torch.scatter(mask, dim=-1, index=temp_target.unsqueeze(-1),
|
||||
value=torch.finfo(student_scores.dtype).min)
|
||||
return loss
|
||||
|
||||
def forward(self, query: Dict[str, Tensor] = None, passage: Dict[str, Tensor] = None, teacher_scores: Tensor = None,
|
||||
bi_directions=None):
|
||||
if self.enable_sub_batch:
|
||||
q_dense_vecs, q_sparse_vecs, q_colbert_vecs = self.encode(query,
|
||||
sub_batch_size=self.compute_sub_batch_size(query))
|
||||
p_dense_vecs, p_sparse_vecs, p_colbert_vecs = self.encode(passage,
|
||||
sub_batch_size=self.compute_sub_batch_size(
|
||||
passage))
|
||||
else:
|
||||
q_dense_vecs, q_sparse_vecs, q_colbert_vecs = self.encode(query)
|
||||
p_dense_vecs, p_sparse_vecs, p_colbert_vecs = self.encode(passage)
|
||||
|
||||
if self.training:
|
||||
if teacher_scores is not None:
|
||||
# print("Use soft-label distillation...")
|
||||
teacher_targets = F.softmax(teacher_scores, dim=-1) # B N
|
||||
group_size = p_dense_vecs.size(0) // q_dense_vecs.size(0)
|
||||
|
||||
# dense loss
|
||||
dense_scores = self.dense_score(q_dense_vecs, p_dense_vecs) # B, B * N
|
||||
if self.negatives_cross_device:
|
||||
cross_q_dense_vecs = self._dist_gather_tensor(q_dense_vecs)
|
||||
cross_p_dense_vecs = self._dist_gather_tensor(p_dense_vecs)
|
||||
cross_teacher_targets = self._dist_gather_tensor(teacher_targets)
|
||||
cross_dense_scores = self.dense_score(cross_q_dense_vecs, cross_p_dense_vecs)
|
||||
|
||||
loss = self.distill_loss(cross_teacher_targets, cross_dense_scores, group_size=group_size)
|
||||
else:
|
||||
loss = self.distill_loss(teacher_targets, dense_scores, group_size=group_size)
|
||||
|
||||
if self.unified_finetuning:
|
||||
# sparse and colbert loss
|
||||
sparse_scores = self.sparse_score(q_sparse_vecs, p_sparse_vecs) # B, B * N
|
||||
sparse_loss = self.distill_loss(teacher_targets, sparse_scores, group_size=group_size)
|
||||
|
||||
colbert_scores = self.colbert_score(q_colbert_vecs, p_colbert_vecs,
|
||||
q_mask=query['attention_mask']) # B, B * N
|
||||
colbert_loss = self.distill_loss(teacher_targets, colbert_scores, group_size=group_size)
|
||||
|
||||
ensemble_loss = self.distill_loss(teacher_targets,
|
||||
dense_scores + 0.3 * sparse_scores + colbert_scores,
|
||||
group_size=group_size)
|
||||
loss = (loss + ensemble_loss + 0.1 * sparse_loss + colbert_loss) / 4
|
||||
|
||||
|
||||
else:
|
||||
idxs = torch.arange(q_dense_vecs.size(0), device=q_dense_vecs.device, dtype=torch.long)
|
||||
targets = idxs * (p_dense_vecs.size(0) // q_dense_vecs.size(0))
|
||||
|
||||
# dense loss
|
||||
dense_scores = self.dense_score(q_dense_vecs, p_dense_vecs) # B, B * N
|
||||
if self.negatives_cross_device:
|
||||
cross_q_dense_vecs = self._dist_gather_tensor(q_dense_vecs)
|
||||
cross_p_dense_vecs = self._dist_gather_tensor(p_dense_vecs)
|
||||
|
||||
cross_idxs = torch.arange(cross_q_dense_vecs.size(0), device=cross_q_dense_vecs.device, dtype=torch.long)
|
||||
|
||||
cross_targets = cross_idxs * (cross_p_dense_vecs.size(0) // cross_q_dense_vecs.size(0))
|
||||
cross_dense_scores = self.dense_score(cross_q_dense_vecs, cross_p_dense_vecs)
|
||||
|
||||
loss = self.compute_loss(cross_dense_scores, cross_targets)
|
||||
else:
|
||||
loss = self.compute_loss(dense_scores, targets)
|
||||
|
||||
if self.unified_finetuning:
|
||||
# sparse and colbert loss
|
||||
sparse_scores = self.sparse_score(q_sparse_vecs, p_sparse_vecs) # B, B * N
|
||||
sparse_loss = self.compute_loss(sparse_scores, targets)
|
||||
|
||||
colbert_scores = self.colbert_score(q_colbert_vecs, p_colbert_vecs,
|
||||
q_mask=query['attention_mask']) # B, B * N
|
||||
colbert_loss = self.compute_loss(colbert_scores, targets)
|
||||
|
||||
ensemble_loss = self.compute_loss(dense_scores + 0.3 * sparse_scores + colbert_scores, targets)
|
||||
loss = (loss + ensemble_loss + 0.1 * sparse_loss + colbert_loss) / 4
|
||||
|
||||
if self.use_self_distill and self.step > self.self_distill_start_step and self.unified_finetuning:
|
||||
ensemble_scores = dense_scores + 0.3 * sparse_scores + colbert_scores
|
||||
teacher_targets = torch.softmax(ensemble_scores.detach(), dim=-1)
|
||||
ensemble_distill_dense_loss = - torch.mean(
|
||||
torch.sum(torch.log_softmax(dense_scores, dim=-1) * teacher_targets, dim=-1))
|
||||
ensemble_distill_sparse_loss = - torch.mean(
|
||||
torch.sum(torch.log_softmax(sparse_scores, dim=-1) * teacher_targets, dim=-1))
|
||||
ensemble_distill_colbert_loss = - torch.mean(
|
||||
torch.sum(torch.log_softmax(colbert_scores, dim=-1) * teacher_targets, dim=-1))
|
||||
loss += (ensemble_distill_dense_loss + 0.1 * ensemble_distill_sparse_loss + ensemble_distill_colbert_loss) / 3
|
||||
loss = loss / 2
|
||||
self.step += 1
|
||||
else:
|
||||
loss = None
|
||||
return EncoderOutput(
|
||||
loss=loss,
|
||||
)
|
||||
|
||||
def compute_loss(self, scores, target):
|
||||
return self.cross_entropy(scores, target)
|
||||
|
||||
def _dist_gather_tensor(self, t: Optional[torch.Tensor]):
|
||||
if t is None:
|
||||
return None
|
||||
t = t.contiguous()
|
||||
|
||||
all_tensors = [torch.empty_like(t) for _ in range(self.world_size)]
|
||||
dist.all_gather(all_tensors, t)
|
||||
|
||||
all_tensors[self.process_rank] = t
|
||||
all_tensors = torch.cat(all_tensors, dim=0)
|
||||
|
||||
return all_tensors
|
||||
|
||||
def save(self, output_dir: str):
|
||||
def _trans_state_dict(state_dict):
|
||||
state_dict = type(state_dict)(
|
||||
{k: v.clone().cpu()
|
||||
for k,
|
||||
v in state_dict.items()})
|
||||
return state_dict
|
||||
|
||||
self.model.save_pretrained(output_dir, state_dict=_trans_state_dict(self.model.state_dict()))
|
||||
|
||||
if self.unified_finetuning:
|
||||
torch.save(_trans_state_dict(self.colbert_linear.state_dict()),
|
||||
os.path.join(output_dir, 'colbert_linear.pt'))
|
||||
torch.save(_trans_state_dict(self.sparse_linear.state_dict()),
|
||||
os.path.join(output_dir, 'sparse_linear.pt'))
|
||||
|
||||
def load_pooler(self, model_dir):
|
||||
colbert_state_dict = torch.load(os.path.join(model_dir, 'colbert_linear.pt'), map_location='cpu')
|
||||
sparse_state_dict = torch.load(os.path.join(model_dir, 'sparse_linear.pt'), map_location='cpu')
|
||||
self.colbert_linear.load_state_dict(colbert_state_dict)
|
||||
self.sparse_linear.load_state_dict(sparse_state_dict)
|
||||
|
||||
|
||||
class BGEM3ForInference(BGEM3Model):
|
||||
|
||||
def forward(self,
|
||||
text_input: Dict[str, Tensor] = None,
|
||||
return_dense: bool = True,
|
||||
return_sparse: bool = False,
|
||||
return_colbert: bool = False,
|
||||
return_sparse_embedding: bool = False):
|
||||
assert return_dense or return_sparse or return_colbert, 'Must choose one or more from `return_colbert`, `return_sparse`, `return_dense` to set `True`!'
|
||||
|
||||
# this is for sparse embedding computation: using optimization suggestion from
|
||||
# issue #1364: https://github.com/FlagOpen/FlagEmbedding/issues/1364
|
||||
self.training = False
|
||||
|
||||
last_hidden_state = self.model(**text_input, return_dict=True).last_hidden_state
|
||||
|
||||
output = {}
|
||||
if return_dense:
|
||||
dense_vecs = self.dense_embedding(last_hidden_state, text_input['attention_mask'])
|
||||
output['dense_vecs'] = dense_vecs
|
||||
if return_sparse:
|
||||
sparse_vecs = self.sparse_embedding(last_hidden_state, text_input['input_ids'],
|
||||
return_embedding=return_sparse_embedding)
|
||||
output['sparse_vecs'] = sparse_vecs
|
||||
if return_colbert:
|
||||
colbert_vecs = self.colbert_embedding(last_hidden_state, text_input['attention_mask'])
|
||||
output['colbert_vecs'] = colbert_vecs
|
||||
|
||||
if self.normlized:
|
||||
if 'dense_vecs' in output:
|
||||
output['dense_vecs'] = torch.nn.functional.normalize(output['dense_vecs'], dim=-1)
|
||||
if 'colbert_vecs' in output:
|
||||
output['colbert_vecs'] = torch.nn.functional.normalize(output['colbert_vecs'], dim=-1)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,155 @@
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
import torch.distributed as dist
|
||||
|
||||
from transformers import AutoConfig, AutoTokenizer
|
||||
from transformers import (
|
||||
HfArgumentParser,
|
||||
set_seed,
|
||||
)
|
||||
from transformers import (
|
||||
TrainerCallback,
|
||||
TrainingArguments,
|
||||
TrainerState,
|
||||
TrainerControl
|
||||
)
|
||||
|
||||
from .arguments import ModelArguments, DataArguments, \
|
||||
RetrieverTrainingArguments as TrainingArguments
|
||||
from .data import SameDatasetTrainDataset, EmbedCollator
|
||||
from .modeling import BGEM3Model
|
||||
from .trainer import BiTrainer
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TrainerCallbackForDataRefresh(TrainerCallback):
|
||||
def __init__(self, train_dataset):
|
||||
self.train_dataset = train_dataset
|
||||
|
||||
def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
|
||||
"""
|
||||
Event called at the end of an epoch.
|
||||
"""
|
||||
self.train_dataset.refresh_epoch()
|
||||
|
||||
|
||||
def main():
|
||||
parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
|
||||
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
|
||||
model_args: ModelArguments
|
||||
data_args: DataArguments
|
||||
training_args: TrainingArguments
|
||||
|
||||
if (
|
||||
os.path.exists(training_args.output_dir)
|
||||
and os.listdir(training_args.output_dir)
|
||||
and training_args.do_train
|
||||
and not training_args.overwrite_output_dir
|
||||
):
|
||||
raise ValueError(
|
||||
f"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome."
|
||||
)
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
|
||||
)
|
||||
logger.warning(
|
||||
"Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
|
||||
training_args.local_rank,
|
||||
training_args.device,
|
||||
training_args.n_gpu,
|
||||
bool(training_args.local_rank != -1),
|
||||
training_args.fp16,
|
||||
)
|
||||
logger.info("Training/evaluation parameters %s", training_args)
|
||||
logger.info("Model parameters %s", model_args)
|
||||
logger.info("Data parameters %s", data_args)
|
||||
|
||||
# Set seed
|
||||
set_seed(training_args.seed)
|
||||
|
||||
num_labels = 1
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
|
||||
cache_dir=model_args.cache_dir,
|
||||
use_fast=False,
|
||||
)
|
||||
config = AutoConfig.from_pretrained(
|
||||
model_args.config_name if model_args.config_name else model_args.model_name_or_path,
|
||||
num_labels=num_labels,
|
||||
cache_dir=model_args.cache_dir,
|
||||
)
|
||||
logger.info('Config: %s', config)
|
||||
|
||||
model = BGEM3Model(model_name=model_args.model_name_or_path,
|
||||
normlized=training_args.normlized,
|
||||
sentence_pooling_method=training_args.sentence_pooling_method,
|
||||
negatives_cross_device=training_args.negatives_cross_device,
|
||||
temperature=training_args.temperature,
|
||||
enable_sub_batch=training_args.enable_sub_batch,
|
||||
unified_finetuning=training_args.unified_finetuning,
|
||||
use_self_distill=training_args.use_self_distill,
|
||||
colbert_dim=training_args.colbert_dim,
|
||||
self_distill_start_step=training_args.self_distill_start_step)
|
||||
|
||||
if training_args.fix_position_embedding:
|
||||
for k, v in model.named_parameters():
|
||||
if "position_embeddings" in k:
|
||||
logging.info(f"Freeze the parameters for {k}")
|
||||
v.requires_grad = False
|
||||
if training_args.fix_encoder:
|
||||
for k, v in model.named_parameters():
|
||||
if "colbert_linear" in k or 'sparse_linear' in k:
|
||||
logging.info(f"train the parameters for {k}")
|
||||
else:
|
||||
v.requires_grad = False
|
||||
|
||||
# print(f"===========================Rank {dist.get_rank()}: start loading data===========================")
|
||||
if data_args.same_task_within_batch:
|
||||
train_dataset = SameDatasetTrainDataset(args=data_args,
|
||||
batch_size=training_args.per_device_train_batch_size,
|
||||
seed=training_args.seed,
|
||||
num_processes=training_args.world_size,
|
||||
process_index=training_args.process_index)
|
||||
training_args.per_device_train_batch_size = 1
|
||||
training_args.dataloader_num_workers = 0 # avoid multi-processes
|
||||
else:
|
||||
raise NotImplementedError("Not support `same_task_within_batch=False`")
|
||||
|
||||
data_collator = EmbedCollator(
|
||||
tokenizer,
|
||||
query_max_len=data_args.query_max_len,
|
||||
passage_max_len=data_args.passage_max_len
|
||||
)
|
||||
|
||||
trainer = BiTrainer(
|
||||
model=model,
|
||||
args=training_args,
|
||||
train_dataset=train_dataset,
|
||||
data_collator=data_collator,
|
||||
tokenizer=tokenizer
|
||||
)
|
||||
|
||||
if data_args.same_task_within_batch:
|
||||
trainer.add_callback(TrainerCallbackForDataRefresh(train_dataset))
|
||||
|
||||
Path(training_args.output_dir).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Training
|
||||
# print(f"===========================Rank {dist.get_rank()}: start training===========================")
|
||||
trainer.train()
|
||||
trainer.save_model()
|
||||
# For convenience, we also re-save the tokenizer to the same directory,
|
||||
# so that you can share your model easily on huggingface.co/models =)
|
||||
if trainer.is_world_process_zero():
|
||||
tokenizer.save_pretrained(training_args.output_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,209 @@
|
||||
"""
|
||||
python split_data_by_length.py \
|
||||
--input_path train_data \
|
||||
--output_dir train_data_split \
|
||||
--cache_dir .cache \
|
||||
--log_name .split_log \
|
||||
--length_list 0 500 1000 2000 3000 4000 5000 6000 7000 \
|
||||
--model_name_or_path BAAI/bge-m3 \
|
||||
--num_proc 16 \
|
||||
--overwrite False
|
||||
"""
|
||||
import os
|
||||
import json
|
||||
import math
|
||||
import time
|
||||
import argparse
|
||||
import datasets
|
||||
from tqdm import tqdm
|
||||
from pprint import pprint
|
||||
from transformers import AutoTokenizer
|
||||
from datasets import load_dataset, Features, Value, Sequence
|
||||
|
||||
|
||||
def get_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--input_path', type=str, required=True, help='the path of input datas')
|
||||
parser.add_argument('--output_dir', type=str, required=True, help='the dir of output datas')
|
||||
parser.add_argument('--cache_dir', type=str, default=None, help='the cache dir')
|
||||
parser.add_argument('--log_name', type=str, default='.split_log', help='the name of log file, default: `.split_log`, which will be saved to `output_dir`')
|
||||
parser.add_argument('--length_list', type=int, default=[0, 500, 1000, 2000, 3000, 4000, 5000, 6000, 7000], nargs='+', help='the length list to split')
|
||||
parser.add_argument('--model_name_or_path', type=str, default='BAAI/bge-m3', help='the model name or path of the tokenizer')
|
||||
parser.add_argument('--num_proc', type=int, default=16, help='the number of process, default: 16')
|
||||
parser.add_argument('--overwrite', action='store_true', default=False, help='whether to overwrite the output file, default: False')
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
class SplitByLengthHandler:
|
||||
def __init__(self,
|
||||
model_name_or_path: str,
|
||||
cache_dir: str=None,
|
||||
num_proc: int=16,
|
||||
length_list: list=[0, 500, 1000, 2000, 3000, 4000, 5000, 6000, 7000],
|
||||
overwrite: bool=False):
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
|
||||
self.cache_dir = cache_dir
|
||||
self.num_proc = num_proc
|
||||
self.length_ranges_list = self._get_length_ranges_list(length_list)
|
||||
self.overwrite = overwrite
|
||||
|
||||
pprint(self.length_ranges_list)
|
||||
|
||||
def _map_func(examples):
|
||||
results = {}
|
||||
results['idx'] = []
|
||||
results['max_length'] = []
|
||||
for i in range(len(examples['query'])):
|
||||
idx = examples['idx'][i]
|
||||
query = examples['query'][i]
|
||||
pos, neg = examples['pos'][i], examples['neg'][i]
|
||||
all_texts = [query] + pos + neg
|
||||
|
||||
max_len = 0
|
||||
for x in all_texts:
|
||||
tokenized_x = self.tokenizer(x)['input_ids']
|
||||
if len(tokenized_x) > max_len:
|
||||
max_len = len(tokenized_x)
|
||||
|
||||
results['idx'].append(idx)
|
||||
results['max_length'].append(max_len)
|
||||
return results
|
||||
|
||||
self._map_func = _map_func
|
||||
|
||||
@staticmethod
|
||||
def _get_length_ranges_list(length_list: list):
|
||||
length_ranges_list = []
|
||||
length_list = sorted(length_list)
|
||||
for i in range(len(length_list)):
|
||||
length_l = length_list[i]
|
||||
if i == len(length_list) - 1:
|
||||
length_r = math.inf
|
||||
else:
|
||||
length_r = length_list[i + 1]
|
||||
assert 0 <= length_l < length_r
|
||||
length_ranges_list.append((length_l, length_r))
|
||||
|
||||
return length_ranges_list
|
||||
|
||||
def _process_dir(self, dir_path: str, output_dir: str):
|
||||
assert os.path.isdir(dir_path)
|
||||
log_info_list = []
|
||||
for file in tqdm(os.listdir(dir_path), desc=f'processing {dir_path}'):
|
||||
file_path = os.path.join(dir_path, file)
|
||||
if not file_path.endswith('.jsonl'):
|
||||
print(f"skip {file_path} ...")
|
||||
continue
|
||||
|
||||
output_path = os.path.join(output_dir, '.'.join(file.split('.')[:-1]))
|
||||
log_info = self._process_file(file_path, output_path)
|
||||
log_info_list.append(log_info)
|
||||
return log_info_list
|
||||
|
||||
def _process_file(self, file_path: str, output_path: str):
|
||||
assert not os.path.isdir(file_path)
|
||||
|
||||
start_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
|
||||
|
||||
features = Features({
|
||||
'query': Value('string'),
|
||||
'pos': Sequence(Value('string')),
|
||||
'neg': Sequence(Value('string'))
|
||||
})
|
||||
kd_features = Features({
|
||||
'query': Value('string'),
|
||||
'pos': Sequence(Value('string')),
|
||||
'neg': Sequence(Value('string')),
|
||||
'pos_scores': Sequence(Value('float')),
|
||||
'neg_scores': Sequence(Value('float'))
|
||||
})
|
||||
try:
|
||||
dataset = load_dataset('json', data_files=file_path, cache_dir=self.cache_dir, features=features)['train']
|
||||
except:
|
||||
dataset = load_dataset('json', data_files=file_path, cache_dir=self.cache_dir, features=kd_features)['train']
|
||||
|
||||
dataset_with_idx_list = []
|
||||
for i, data in enumerate(dataset):
|
||||
data['idx'] = i
|
||||
dataset_with_idx_list.append(data)
|
||||
dataset_with_idx = datasets.Dataset.from_list(dataset_with_idx_list)
|
||||
|
||||
mapped_dataset = dataset_with_idx.map(self._map_func, batched=True, num_proc=self.num_proc)
|
||||
|
||||
split_info_dict = {}
|
||||
for length_l, length_r in self.length_ranges_list:
|
||||
save_path = output_path + f'_len-{length_l}-{length_r}.jsonl'
|
||||
if os.path.exists(save_path) and not self.overwrite:
|
||||
print(f'{save_path} exists, skip')
|
||||
continue
|
||||
|
||||
idxs = mapped_dataset.filter(lambda x: length_l <= x['max_length'] < length_r, num_proc=self.num_proc)
|
||||
split_dataset = dataset_with_idx.select(idxs['idx'])
|
||||
split_dataset = split_dataset.remove_columns('idx')
|
||||
|
||||
split_info_dict[f'len-{length_l}-{length_r}'] = len(split_dataset)
|
||||
|
||||
if len(split_dataset) > 0:
|
||||
split_dataset.to_json(save_path, force_ascii=False)
|
||||
|
||||
end_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
|
||||
|
||||
size = len(dataset)
|
||||
avg_length = sum(mapped_dataset['max_length']) / size
|
||||
log_info = {
|
||||
'file_name': os.path.basename(file_path),
|
||||
'size': size,
|
||||
'avg_length': avg_length,
|
||||
'file_path': file_path,
|
||||
'start_time': start_time,
|
||||
'end_time': end_time,
|
||||
'split_info': split_info_dict
|
||||
}
|
||||
return log_info
|
||||
|
||||
def run(self, input_path: str, output_dir: str, log_name: str=None):
|
||||
if not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
|
||||
if log_name is None:
|
||||
log_path = os.path.join(output_dir, '.split_log')
|
||||
else:
|
||||
log_path = os.path.join(output_dir, log_name)
|
||||
|
||||
log_info_list = []
|
||||
|
||||
if os.path.isdir(input_path):
|
||||
log_info_list = self._process_dir(input_path, output_dir)
|
||||
else:
|
||||
file_name = os.path.basename(input_path)
|
||||
output_path = os.path.join(output_dir, '.'.join(file_name.split('.')[:-1]))
|
||||
log_info = self._process_file(input_path, output_path)
|
||||
log_info_list.append(log_info)
|
||||
|
||||
with open(log_path, 'a', encoding='utf-8') as f:
|
||||
for log_info in log_info_list:
|
||||
json.dump(log_info, f, ensure_ascii=False)
|
||||
f.write('\n')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = get_args()
|
||||
input_path = args.input_path
|
||||
output_dir = args.output_dir
|
||||
log_name = args.log_name
|
||||
|
||||
handler = SplitByLengthHandler(
|
||||
model_name_or_path=args.model_name_or_path,
|
||||
cache_dir=args.cache_dir,
|
||||
num_proc=args.num_proc,
|
||||
length_list=args.length_list if isinstance(args.length_list, list) else [args.length_list],
|
||||
overwrite=args.overwrite
|
||||
)
|
||||
|
||||
handler.run(
|
||||
input_path=input_path,
|
||||
output_dir=output_dir,
|
||||
log_name=log_name
|
||||
)
|
||||
print('\nDONE!')
|
||||
@@ -0,0 +1,51 @@
|
||||
from sentence_transformers import SentenceTransformer, models
|
||||
from transformers.trainer import *
|
||||
|
||||
|
||||
def save_ckpt_for_sentence_transformers(ckpt_dir, pooling_mode: str = 'cls', normlized: bool=True):
|
||||
word_embedding_model = models.Transformer(ckpt_dir)
|
||||
pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), pooling_mode=pooling_mode)
|
||||
if normlized:
|
||||
normlize_layer = models.Normalize()
|
||||
model = SentenceTransformer(modules=[word_embedding_model, pooling_model, normlize_layer], device='cpu')
|
||||
else:
|
||||
model = SentenceTransformer(modules=[word_embedding_model, pooling_model], device='cpu')
|
||||
model.save(ckpt_dir)
|
||||
|
||||
|
||||
class BiTrainer(Trainer):
|
||||
def _save(self, output_dir: Optional[str] = None, state_dict=None):
|
||||
output_dir = output_dir if output_dir is not None else self.args.output_dir
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
logger.info("Saving model checkpoint to %s", output_dir)
|
||||
# Save a trained model and configuration using `save_pretrained()`.
|
||||
# They can then be reloaded using `from_pretrained()`
|
||||
if not hasattr(self.model, 'save'):
|
||||
raise NotImplementedError(
|
||||
f'MODEL {self.model.__class__.__name__} '
|
||||
f'does not support save interface')
|
||||
else:
|
||||
self.model.save(output_dir)
|
||||
if self.tokenizer is not None and self.is_world_process_zero():
|
||||
self.tokenizer.save_pretrained(output_dir)
|
||||
|
||||
torch.save(self.args, os.path.join(output_dir, "training_args.bin"))
|
||||
|
||||
# save the checkpoint for sentence-transformers library
|
||||
if self.is_world_process_zero():
|
||||
save_ckpt_for_sentence_transformers(output_dir,
|
||||
pooling_mode=self.args.sentence_pooling_method,
|
||||
normlized=self.args.normlized)
|
||||
|
||||
|
||||
def compute_loss(self, model, inputs, return_outputs=False):
|
||||
"""
|
||||
How the loss is computed by Trainer. By default, all models return the loss in the first element.
|
||||
|
||||
Subclass and override for custom behavior.
|
||||
"""
|
||||
|
||||
outputs = model(**inputs)
|
||||
loss = outputs.loss
|
||||
|
||||
return (loss, outputs) if return_outputs else loss
|
||||
@@ -0,0 +1,149 @@
|
||||
<div align="center">
|
||||
<h1> BGE-Reasoner: Towards End-to-End Reasoning-Intensive Information Retrieval </h1>
|
||||
</div>
|
||||
|
||||
## Introduction
|
||||
|
||||
We introduce **BGE-Reasoner**, an end-to-end reasoning-intensive information retrieval framework. BGE-Reasoner is characterized by three key features:
|
||||
|
||||
1. **End-to-end**: It comprises three core components in IR—**BGE-Reasoner-Rewriter**, **BGE-Reasoner-Embed**, and **BGE-Reasoner-Reranker**—covering the entire retrieval pipeline, from query rewriting and retrieval to reranking for reasoning-intensive tasks.
|
||||
2. **Excellent performance**: **BGE-Reasoner** achieves **state-of-the-art (SOTA)** performance on [BRIGHT](https://brightbenchmark.github.io/), a reasoning-intensive information retrieval benchmark, with an **nDCG@10 of 45.2** across 12 datasets (released on Aug 21, 2025), outperforming the previous SOTA by +3.6 points (41.6 from [DIVER](https://arxiv.org/pdf/2508.07995), Aug 12, 2025).
|
||||
3. **Open-source resources**: We will release the code, model checkpoints, training data, and evaluation scripts to facilitate future research on reasoning-intensive information retrieval. Please stay tuned!
|
||||
|
||||
|
||||
## Open-source resources
|
||||
|
||||
| Resource Type | Name | Link | Release Date | Comments |
|
||||
| ------------------ | --------------------- | ----------- | ------------------ | ------------------ |
|
||||
| Model | BGE-Reasoner-Rewriter | [🤗]() (TBA) | - | |
|
||||
| Model | BGE-Reasoner-Reranker | [🤗]() (TBA) | - | |
|
||||
| Model | BGE-Reasoner-Embed-Qwen3-8B-0923 | [🤗](https://huggingface.co/BAAI/bge-reasoner-embed-qwen3-8b-0923) | Sep 23, 2025 | nDCG@10 = 37.1 using original query, fine-tuned on [Qwen/Qwen3-8B](https://huggingface.co/Qwen/Qwen3-8B) with our latest refined training data (data to be released) |
|
||||
| Search Results | BGE-Reasoner-Embed-Qwen3-8B-0923 Search Results | [🤗](https://huggingface.co/BAAI/bge-reasoner-embed-qwen3-8b-0923/tree/main/search_results) | Sep 23, 2025 | nDCG@10 = 37.1 using original query |
|
||||
| Search Results | BGE-Reasoner-Embed-0821 Search Results | [🤗](https://huggingface.co/datasets/hanhainebula/bright-search-results_bge-reasoner-embed-0821/tree/main) | Sep 4, 2025 | nDCG@10 = 32.5 using original query, submission to BRIGHT leaderboard on Aug 21, 2025 |
|
||||
| Training Data | BGE-Reasoner-Data | [🤗](https://huggingface.co/datasets/hanhainebula/bge-reasoner-data/tree/main/bge-reasoner-data-0904) | Sep 4, 2025 | part of our training data; full data to be released in the future |
|
||||
| Evaluation Scripts | - | (TBA) | - | |
|
||||
|
||||
|
||||
|
||||
## Performance
|
||||
|
||||
**BGE-Reasoner** achieves SOTA performance on the **BRIGHT** benchmark with the following pipeline:
|
||||
|
||||
|
||||
|
||||
|
||||
1. **Query Rewrite**: **BGE-Reasoner-Rewriter** generates 5 rewritten queries for each original query; all 5 rewrites are used for retrieval.
|
||||
2. **Retrieval**: For each rewritten query, **BGE-Reasoner-Embed** and **BM25** retrieve the top-2000 documents. We aggregate results across the 5 rewrites by summing the corresponding scores to produce a final score per method.
|
||||
3. **Reranking**:
|
||||
- We rerank the top-100 documents from each retrieval method using **BGE-Reasoner-Reranker** (models: 8B, 14B, 32B), producing 6 reranked top-10 lists (2 retrieval methods × 3 reranker sizes).
|
||||
- We also create a hybrid top-10 by fusing **BGE-Reasoner-Embed** and **BM25** (weights: 0.75 / 0.25 after min–max normalization).
|
||||
- Finally, we combine the 7 top-10 lists (6 reranked + 1 hybrid) to produce the final top-10.
|
||||
|
||||
|
||||
### Full Pipeline Results
|
||||
|
||||
|
||||
|
||||
Note:
|
||||
- "**Avg - ALL**" refers to the average performance across **all 12 datasets** in the BRIGHT benchmark.
|
||||
- "**Avg - SE**" refers to the average performance across the **7 datasets in the StackExchange subset** of the BRIGHT benchmark.
|
||||
- "**Avg - CD**" refers to the average performance across the **2 datasets in the Coding subset** of the BRIGHT benchmark.
|
||||
- "**Avg - MT**" refers to the average performance across the **3 datasets in the Theorem-based subset** of the BRIGHT benchmark.
|
||||
|
||||
> Sources of results:
|
||||
>
|
||||
> [1] https://arxiv.org/pdf/2504.20595
|
||||
>
|
||||
> [2] https://github.com/Debrup-61/RaDeR
|
||||
>
|
||||
> [3] https://huggingface.co/ielabgroup/Rank-R1-32B-v0.2
|
||||
>
|
||||
> [4] https://github.com/jataware/XRR2
|
||||
>
|
||||
> [5] http://arxiv.org/pdf/2508.07050
|
||||
>
|
||||
> [6] https://arxiv.org/pdf/2508.07995
|
||||
|
||||
|
||||
### Embedder & Rewriter Results
|
||||
|
||||
|
||||
#### BGE-Reasoner-Embed-Qwen3-8B-0923
|
||||
|
||||
**BGE-Reasoner-Embed-Qwen3-8B-0923**, fine-tuned on [Qwen/Qwen3-8B](https://huggingface.co/Qwen/Qwen3-8B) with our latest refined training data (data to be released), achieves strong performance on the BRIGHT benchmark:
|
||||
|
||||
- With original queries, it attains **nDCG@10 = 37.1**, an absolute improvement of **+8.2** over the previous best ([DIVER](https://arxiv.org/pdf/2508.07995): 28.9).
|
||||
- Using the GPT-4 reasoning queries provided by BRIGHT, the score increases to **39.7**, which is **+7.6** higher than DIVER’s corresponding result (32.1).
|
||||
|
||||
> On Sep 23, 2025, we released the first-stage search results of BGE-Reasoner-Embed-Qwen3-8B-0923 using original queries and GPT-4 reasoning queries (Top-2000 candidates; excluded IDs removed) [here](https://huggingface.co/BAAI/bge-reasoner-embed-qwen3-8b-0923/tree/main/search_results). The model checkpoint is available [here](https://huggingface.co/BAAI/bge-reasoner-embed-qwen3-8b-0923).
|
||||
|
||||
|
||||
|
||||
Note:
|
||||
- "**Avg - ALL**" refers to the average performance across **all 12 datasets** in the BRIGHT benchmark.
|
||||
- "**Avg - SE**" refers to the average performance across the **7 datasets in the StackExchange subset** of the BRIGHT benchmark.
|
||||
- "**Avg - CD**" refers to the average performance across the **2 datasets in the Coding subset** of the BRIGHT benchmark.
|
||||
- "**Avg - MT**" refers to the average performance across the **3 datasets in the Theorem-based subset** of the BRIGHT benchmark.
|
||||
|
||||
> Sources of Results:
|
||||
>
|
||||
> [1] https://arxiv.org/pdf/2407.12883
|
||||
>
|
||||
> [2] https://arxiv.org/pdf/2504.20595
|
||||
>
|
||||
> [3] https://github.com/Debrup-61/RaDeR
|
||||
>
|
||||
> [4] https://seed1-5-embedding.github.io
|
||||
>
|
||||
> [5] https://arxiv.org/pdf/2508.07995
|
||||
>
|
||||
> *: results evaluated with our script
|
||||
|
||||
#### BGE-Reasoner-Embed-0821
|
||||
|
||||
**BGE-Reasoner-Embed-0821**, submitted to the BRIGHT leaderboard on Aug 21, 2025, achieves excellent performance on the benchmark:
|
||||
|
||||
- With original queries, it attains **nDCG@10 = 32.5**, an absolute improvement of **+3.6** over the previous best ([DIVER](https://arxiv.org/pdf/2508.07995): 28.9).
|
||||
- Using the GPT-4 reasoning queries provided by BRIGHT, the score increases to **37.7**, which is **+5.6** higher than DIVER’s corresponding result (32.1). Combining our embedding-based retrieval with BM25 (hybrid fusion, weights: 0.75 / 0.25) yields **nDCG@10 = 40.2**.
|
||||
- Finally, when using rewritten queries produced by **BGE-Reasoner-Rewriter** and fusing with BM25 (weights: 0.75 / 0.25), we reach **nDCG@10 = 40.8**.
|
||||
|
||||
> On Sep 4, 2025, we released the first-stage search results of BGE-Reasoner-Embed-0821 using original queries and GPT-4 reasoning queries (Top-2000 candidates; excluded IDs removed) [here](https://huggingface.co/datasets/hanhainebula/bright-search-results_bge-reasoner-embed-0821/tree/main). The model checkpoint will not be released due to its suboptimal performance compared to BGE-Reasoner-Embed-Qwen3-8B-0923.
|
||||
|
||||
|
||||
|
||||
|
||||
Note:
|
||||
- "**Avg - ALL**" refers to the average performance across **all 12 datasets** in the BRIGHT benchmark.
|
||||
- "**Avg - SE**" refers to the average performance across the **7 datasets in the StackExchange subset** of the BRIGHT benchmark.
|
||||
- "**Avg - CD**" refers to the average performance across the **2 datasets in the Coding subset** of the BRIGHT benchmark.
|
||||
- "**Avg - MT**" refers to the average performance across the **3 datasets in the Theorem-based subset** of the BRIGHT benchmark.
|
||||
|
||||
> Sources of Results:
|
||||
>
|
||||
> [1] https://arxiv.org/pdf/2407.12883
|
||||
>
|
||||
> [2] https://arxiv.org/pdf/2504.20595
|
||||
>
|
||||
> [3] https://github.com/Debrup-61/RaDeR
|
||||
>
|
||||
> [4] https://seed1-5-embedding.github.io
|
||||
>
|
||||
> [5] https://arxiv.org/pdf/2508.07995
|
||||
>
|
||||
> *: results evaluated with our script
|
||||
|
||||
|
||||
## Technical Details
|
||||
|
||||
|
||||
The technical details for each component of **BGE-Reasoner** will be released soon. Please stay tuned!
|
||||
|
||||
|
||||
## Contact Information
|
||||
|
||||
Some resources are not yet publicly available. If you have urgent research needs for any of these resources (e.g., model checkpoints, search results, evaluation scripts) or have any questions, please contact Jianlyu Chen at jianlvchen@gmail.com.
|
||||
|
||||
|
||||
## Citation
|
||||
|
||||
TBA
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 143 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 133 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 174 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 58 KiB |
@@ -0,0 +1,22 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2024 JUNJIE99
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
|
||||
@@ -0,0 +1,198 @@
|
||||
<h1 align="center">MegaPairs: Massive Data Synthesis For Universal Multimodal Retrieval</h1>
|
||||
|
||||
<p align="center">
|
||||
<a href="https://arxiv.org/abs/2412.14475">
|
||||
<img alt="Build" src="http://img.shields.io/badge/cs.CV-arXiv%3A2412.14475-B31B1B.svg">
|
||||
</a>
|
||||
<a href="https://github.com/VectorSpaceLab/MegaPairs">
|
||||
<img alt="Build" src="https://img.shields.io/badge/Github-Code-blue">
|
||||
</a>
|
||||
<a href="https://huggingface.co/datasets/JUNJIE99/MegaPairs">
|
||||
<img alt="Build" src="https://img.shields.io/badge/🤗 Datasets-MegaPairs-yellow">
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
</a>
|
||||
<a href="https://huggingface.co/BAAI/BGE-VL-base">
|
||||
<img alt="Build" src="https://img.shields.io/badge/🤗 Model-BGE_VL_base-yellow">
|
||||
</a>
|
||||
<a href="https://huggingface.co/BAAI/BGE-VL-large">
|
||||
<img alt="Build" src="https://img.shields.io/badge/🤗 Model-BGE_VL_large-yellow">
|
||||
</a>
|
||||
<a href="https://huggingface.co/BAAI/BGE-VL-MLLM-S1">
|
||||
<img alt="Build" src="https://img.shields.io/badge/🤗 Model-BGE_VL_MLLM_S1-yellow">
|
||||
</a>
|
||||
<a href="https://huggingface.co/BAAI/BGE-VL-MLLM-S2">
|
||||
<img alt="Build" src="https://img.shields.io/badge/🤗 Model-BGE_VL_MLLM_S2-yellow">
|
||||
</a>
|
||||
</p>
|
||||
<p align="center">
|
||||
|
||||
|
||||
## News
|
||||
```2025-4-13``` 🎉🎉 We have uploaded our MegaPairs dataset to [🤗Hugging Face](https://huggingface.co/datasets/JUNJIE99/MegaPairs), which contains over 26 million multimodal retrieval instruction-tuning triplets. To reduce upload time and enhance data accessibility, we resized all images to a resolution of 512 × 512 instead of using their original size. This adjustment has minimal impact on performance, considering that most vision-language models (e.g., CLIP) use even smaller input image sizes. [Dataset Card](https://github.com/VectorSpaceLab/MegaPairs?tab=readme-ov-file#megapairs-dataset-card)
|
||||
|
||||
```2025-4-2``` 🌟🌟 BGE-VL models are also available on [WiseModel](https://www.wisemodel.cn/models/JUNJIE99/BGE-VL-large).
|
||||
|
||||
```2025-3-6``` 📰📰 Thank you to [SyncedTech (机器之心)](https://mp.weixin.qq.com/s/iw9BmSDwv6NYtD7pkC5kxQ), [QbitAI (量子位)](https://mp.weixin.qq.com/s/r_zWAZ0ir5732OfIrEsDtg), and [AI Era (新智元)](https://mp.weixin.qq.com/s/FZwKYJnx_78YDAEreu1edg) for reporting on our work!
|
||||
|
||||
```2025-3-4``` 🚀🚀 We have released the BGE-VL-MLLM models on Huggingface: [BGE-VL-MLLM-S1](https://huggingface.co/BAAI/BGE-VL-MLLM-S1) and [BGE-VL-MLLM-S2](https://huggingface.co/BAAI/BGE-VL-MLLM-S2). **BGE-VL-MLLM-S1** is trained exclusively on our MegaPairs dataset, achieving outstanding performance in composed image retrieval, with an 8.1% improvement on the CIRCO benchmark (mAP@5) over the previous state-of-the-art. **BGE-VL-MLLM-S2** builds on BGE-VL-MLLM-S1 with an additional epoch of fine-tuning on the MMEB benchmark training set, delivering enhanced performance across a broader range of multimodal embedding tasks.
|
||||
|
||||
```2024-12-27``` 🚀🚀 BGE-VL-CLIP models are released on Huggingface: [BGE-VL-base](https://huggingface.co/BAAI/BGE-VL-base) and [BGE-VL-large](https://huggingface.co/BAAI/BGE-VL-large).
|
||||
|
||||
```2024-12-19``` 🎉🎉 Release our paper: [MegaPairs: Massive Data Synthesis For Universal Multimodal Retrieval](https://arxiv.org/pdf/2412.14475).
|
||||
|
||||
## Release Plan
|
||||
- [x] Paper
|
||||
- [x] BGE-VL-base and BGE-VL-large models
|
||||
- [x] BGE-VL-MLLM model
|
||||
- [x] MegaPairs Dataset
|
||||
- [x] Evaluation code examples
|
||||
- [ ] Fine-tuning code
|
||||
|
||||
|
||||
## Introduction
|
||||
In this work, we introduce **MegaPairs**, a novel data synthesis method that leverages open-domain images to create *heterogeneous KNN triplets* for universal multimodal retrieval. Our MegaPairs dataset contains over 26 million triplets, and we have trained a series of multimodal retrieval models, **BGE-VL**, including BGE-VL-CLIP (base and large) and BGE-VL-MLLM.
|
||||
|
||||
BGE-VL achieve state-of-the-art performance on four popular zero-shot composed image retrieval benchmarks and the massive multimodal embedding benchmark (MMEB). Extensive experiments demonstrate the ***efficiency, scalability, and generalization*** features of MegaPairs. Please refer to our [paper](https://arxiv.org/abs/2412.14475) for more details.
|
||||
|
||||
|
||||
|
||||
## Model Usage
|
||||
|
||||
### 1. BGE-VL-CLIP Models
|
||||
You can easily use BGE-VL-CLIP models based on ```transformers```
|
||||
> Our code works well on transformers==4.45.2, and we recommend using this version.
|
||||
```python
|
||||
import torch
|
||||
from transformers import AutoModel
|
||||
|
||||
MODEL_NAME = "BAAI/BGE-VL-base" # or "BAAI/BGE-VL-large"
|
||||
|
||||
model = AutoModel.from_pretrained(MODEL_NAME, trust_remote_code=True) # You must set trust_remote_code=True
|
||||
model.set_processor(MODEL_NAME)
|
||||
model.eval()
|
||||
|
||||
with torch.no_grad():
|
||||
query = model.encode(
|
||||
images = "./assets/cir_query.png",
|
||||
text = "Make the background dark, as if the camera has taken the photo at night"
|
||||
)
|
||||
|
||||
candidates = model.encode(
|
||||
images = ["./assets/cir_candi_1.png", "./assets/cir_candi_2.png"]
|
||||
)
|
||||
|
||||
scores = query @ candidates.T
|
||||
print(scores)
|
||||
```
|
||||
|
||||
See the [demo](./retrieval_demo.ipynb) for a complete example of using BGE-VL for multimodel retrieval.
|
||||
|
||||
|
||||
### 2. BGE-VL-MLLM Models
|
||||
|
||||
> Our code works well on transformers==4.45.2, and we recommend using this version.
|
||||
|
||||
```python
|
||||
import torch
|
||||
from transformers import AutoModel
|
||||
from PIL import Image
|
||||
|
||||
MODEL_NAME= "BAAI/BGE-VL-MLLM-S1"
|
||||
|
||||
model = AutoModel.from_pretrained(MODEL_NAME, trust_remote_code=True)
|
||||
model.eval()
|
||||
model.cuda()
|
||||
|
||||
with torch.no_grad():
|
||||
model.set_processor(MODEL_NAME)
|
||||
|
||||
query_inputs = model.data_process(
|
||||
text="Make the background dark, as if the camera has taken the photo at night",
|
||||
images="./assets/cir_query.png",
|
||||
q_or_c="q",
|
||||
task_instruction="Retrieve the target image that best meets the combined criteria by using both the provided image and the image retrieval instructions: "
|
||||
)
|
||||
|
||||
candidate_inputs = model.data_process(
|
||||
images=["./assets/cir_candi_1.png", "./assets/cir_candi_2.png"],
|
||||
q_or_c="c",
|
||||
)
|
||||
|
||||
query_embs = model(**query_inputs, output_hidden_states=True)[:, -1, :]
|
||||
candi_embs = model(**candidate_inputs, output_hidden_states=True)[:, -1, :]
|
||||
|
||||
query_embs = torch.nn.functional.normalize(query_embs, dim=-1)
|
||||
candi_embs = torch.nn.functional.normalize(candi_embs, dim=-1)
|
||||
|
||||
scores = torch.matmul(query_embs, candi_embs.T)
|
||||
print(scores)
|
||||
```
|
||||
|
||||
## MegaPairs Dataset Card
|
||||
|
||||
We are excited to release the **MegaPairs** dataset on [Hugging Face](https://huggingface.co/datasets/JUNJIE99/MegaPairs), which contains over **26 million training samples** tailored for composed image retrieval and universal multimodal retrieval tasks.
|
||||
|
||||
### Dataset Structure
|
||||
|
||||
Each entry in the dataset consists of the following fields:
|
||||
|
||||
- **q_img**: `str`
|
||||
The file path to the query image.
|
||||
|
||||
- **q_text**: `list`
|
||||
A list of textual query statements related to the query image. During training, you can randomly select one statement from this list.
|
||||
|
||||
- **t_img**: `str`
|
||||
The file path to the target image, which serves as the **positive example** for the combination of `q_img` and `q_text`.
|
||||
|
||||
- **hns**: `list`
|
||||
A list of file paths for **hard negative sample** images. These are challenging distractors that are visually or semantically similar to the query. It is recommended to include at least one hard negative sample during training, with **`hns[0]` (the query image itself)** being a mandatory choice. In our experiments, we used **four hard negative samples** per query.
|
||||
|
||||
|
||||
### Usage
|
||||
|
||||
The dataset is available for download and exploration on [Hugging Face](https://huggingface.co/datasets/JUNJIE99/MegaPairs). We encourage researchers and practitioners to leverage this dataset to advance multimodal retrieval research and systems.
|
||||
|
||||
## Model Performance
|
||||
### Zero-Shot Composed Image Retrieval
|
||||
|
||||
BGE-VL sets a new performance benchmark in zero-shot composed image retrieval tasks. On the CIRCO benchmark, our BGE-VL-base model, with only 149 million parameters, surpasses all previous models, including those with 50 times more parameters. Additionally, BGE-VL-MLLM achieves an 8.1% improvement over the previous state-of-the-art model.
|
||||
|
||||
<img src="./assets/res-zs-cir.png" width="800">
|
||||
|
||||
### Zero-Shot Performance on MMEB
|
||||
|
||||
BGE-VL-MLLM achieves state-of-the-art zero-shot performance on the Massive Multimodal Embedding Benchmark (MMEB), despite being trained only on the ImageText-to-Image paradigm. This demonstrates the excellent generalization capability of MegaPairs for multimodal embedding.
|
||||
|
||||
<img src="./assets/res-zs-mmeb.png" width="800">
|
||||
|
||||
### Fine-Tuning Performance on MMEB
|
||||
|
||||
After fine-tuning on downstream tasks, BGE-VL-MLLM maintains its leading performance. Notably, it surpasses the previous state-of-the-art by 7.1% on the MMEB out-of-distribution (OOD) set. These results demonstrate the robust generalization capability of BGE-VL-MLLM and highlight the potential of MegaPairs as foundational training data for universal multimodal embedding.
|
||||
|
||||
<img src="./assets/res-ft-mmeb.png" width="800">
|
||||
|
||||
### Performance Scaling
|
||||
MegaPairs showcases **scalability**: BGE-VL-base improves as training data increases. It also demonstrates **efficiency**: with just 0.5M training samples, BGE-VL-base significantly outperforms MagicLens, which uses the same CLIP-base backbone and was trained on 36.7M samples.
|
||||
|
||||
<img src="./assets/res-scaling.png" width="800">
|
||||
|
||||
|
||||
## License
|
||||
The annotations for MegaPairs and the BGE-VL models are released under the [MIT License](LICENSE). The images in MegaPairs originate from the [Recap-Datacomp](https://huggingface.co/datasets/UCSC-VLAA/Recap-DataComp-1B), which is released under the CC BY 4.0 license.
|
||||
|
||||
|
||||
|
||||
## Citation
|
||||
If you find this repository useful, please consider giving a star ⭐ and citation
|
||||
|
||||
```
|
||||
@article{zhou2024megapairs,
|
||||
title={MegaPairs: Massive Data Synthesis For Universal Multimodal Retrieval},
|
||||
author={Zhou, Junjie and Liu, Zheng and Liu, Ze and Xiao, Shitao and Wang, Yueze and Zhao, Bo and Zhang, Chen Jason and Lian, Defu and Xiong, Yongping},
|
||||
journal={arXiv preprint arXiv:2412.14475},
|
||||
year={2024}
|
||||
}
|
||||
```
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 50 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 880 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 149 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 266 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 142 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 158 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 123 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 181 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 145 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 264 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 178 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 142 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 260 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 145 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 174 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 114 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 133 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 254 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 133 KiB |
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,800 @@
|
||||
{"id": 0, "q_img": "000000281438.jpg", "q_text": "Find a picture that also a man sitting on an outdoor toilet, but has a higher quality and is taken during the daytime."}
|
||||
{"id": 1, "q_img": "000000251485.jpg", "q_text": "Find a picture that also a double bed, but is in a room with wooden walls."}
|
||||
{"id": 2, "q_img": "000000232621.jpg", "q_text": "Find a picture that also a person looking at his laptop, but is sitting on the ground."}
|
||||
{"id": 3, "q_img": "000000258030.jpg", "q_text": "Find a picture that also a man wearing a tie, a shirt and a hat, but is wearing sunglasses."}
|
||||
{"id": 4, "q_img": "000000114366.jpg", "q_text": "Find a picture that also a box of donuts, but shows only one box placed inside a car."}
|
||||
{"id": 5, "q_img": "000000019155.jpg", "q_text": "Find a picture that also an open refrigerator filled with mostly alcoholic beverages, but shows a person in the foreground."}
|
||||
{"id": 6, "q_img": "000000023567.jpg", "q_text": "Find a picture that also a child wearing a chef's hat and an apron, but has more colorful apron and he is looking at the camera."}
|
||||
{"id": 7, "q_img": "000000400924.jpg", "q_text": "Find a picture that also a clock in a big hall, but is taken from the same angle and shows a spiral staircase."}
|
||||
{"id": 8, "q_img": "000000455504.jpg", "q_text": "Find a picture that also a glass bottle on a bench, but is only one and is seen from the side."}
|
||||
{"id": 9, "q_img": "000000240265.jpg", "q_text": "Find a picture that also a double bed, but has blue blankets and a CRT TV next to it."}
|
||||
{"id": 10, "q_img": "000000433496.jpg", "q_text": "Find a picture that also a train wagon seen from the inside, but shows the dining wagon."}
|
||||
{"id": 11, "q_img": "000000549870.jpg", "q_text": "Find a picture that also a person talking on the phone with water in the background, but has only one person and he is lying down."}
|
||||
{"id": 12, "q_img": "000000451179.jpg", "q_text": "Find a picture that also a bowl of fruit in the foreground, but has bananas sliced in half, apples, more fruit and no peach."}
|
||||
{"id": 13, "q_img": "000000539540.jpg", "q_text": "Find a picture that also a sliced carrot, but is on a wooden cutting board with a different arrangement."}
|
||||
{"id": 14, "q_img": "000000067430.jpg", "q_text": "Find a picture that also a canal with boats seen from a river view, but is shot from the same angle and has buildings on both the sides."}
|
||||
{"id": 15, "q_img": "000000117864.jpg", "q_text": "Find a picture that also two people looking at their phones, but are dressed more formally and the photo is taken indoors."}
|
||||
{"id": 16, "q_img": "000000061617.jpg", "q_text": "Find a picture that also a laptop on a desk, but is shot from the side and has several clothes in the background."}
|
||||
{"id": 17, "q_img": "000000119135.jpg", "q_text": "Find a picture that also a living room, but has multiple moving boxes."}
|
||||
{"id": 18, "q_img": "000000490966.jpg", "q_text": "Find a picture that also a kitchen, but is mainly blue and the photo is taken from a different angle."}
|
||||
{"id": 19, "q_img": "000000090604.jpg", "q_text": "Find a picture that also a dog with a Santa hat, but is a statue and the photo is set outdoors."}
|
||||
{"id": 20, "q_img": "000000158745.jpg", "q_text": "Find a picture that also some kites in the sky, but has a bridge in the background."}
|
||||
{"id": 21, "q_img": "000000210061.jpg", "q_text": "Find a picture that also a person seen from the back in greyscale, but is holding an umbrella instead of a suitcase and the photo is taken from the same angle and has no other people in the foreground."}
|
||||
{"id": 22, "q_img": "000000069753.jpg", "q_text": "Find a picture that also a person riding a horse, but is half immersed in the water and no buildings are visible in the background."}
|
||||
{"id": 23, "q_img": "000000281885.jpg", "q_text": "Find a picture that also a person seen from the front sitting on a bench using a laptop, but is shot from the same angle and shows only one person."}
|
||||
{"id": 24, "q_img": "000000152478.jpg", "q_text": "Find a picture that also a birthday cake, but has a bus instead of a horse."}
|
||||
{"id": 25, "q_img": "000000471574.jpg", "q_text": "Find a picture that also a dining table, but is not set and has a clear vase with more white flowers in it."}
|
||||
{"id": 26, "q_img": "000000330427.jpg", "q_text": "Find a picture that also a fire hydrant surrounded by snow, but is of a different color and the photo shows a traffic cone."}
|
||||
{"id": 27, "q_img": "000000032201.jpg", "q_text": "Find a picture that also a person sitting on a bench, but shows fewer people and one parked motor-scooter."}
|
||||
{"id": 28, "q_img": "000000472616.jpg", "q_text": "Find a picture that also a woman taking a selfie in a bathroom mirror, but has only one person and she is wearing a black t-shirt."}
|
||||
{"id": 29, "q_img": "000000094878.jpg", "q_text": "Find a picture that also an abandoned bike, but is in the water and the photo is zoomed out."}
|
||||
{"id": 30, "q_img": "000000194017.jpg", "q_text": "Find a picture that also people riding horses, but shows three of them and there is a forest in the background."}
|
||||
{"id": 31, "q_img": "000000438039.jpg", "q_text": "Find a picture that also a woman on a bench, but is shot in greyscale, the photo has a different background and the woman is looking at a phone."}
|
||||
{"id": 32, "q_img": "000000574169.jpg", "q_text": "Find a picture that also a dog with a book in front of him, but is of a different breed and the book is open."}
|
||||
{"id": 33, "q_img": "000000510439.jpg", "q_text": "Find a picture that also a small bird, but shows only one of them and it is on a metal railing instead of power lines."}
|
||||
{"id": 34, "q_img": "000000021141.jpg", "q_text": "Find a picture that also a white and grey bird, but has the same color and is standing still on the shoreline."}
|
||||
{"id": 35, "q_img": "000000157409.jpg", "q_text": "Find a picture that also a kitchen, but is blue and has more objects attached to the walls."}
|
||||
{"id": 36, "q_img": "000000020668.jpg", "q_text": "Find a picture that also a white Wii controller, but is only one and it is touched by a cat."}
|
||||
{"id": 37, "q_img": "000000066227.jpg", "q_text": "Find a picture that also a close-up open refrigerator seen from the front, but is more stocked and the photo has one person in the foreground."}
|
||||
{"id": 38, "q_img": "000000430952.jpg", "q_text": "Find a picture that also an elephant, but is crossing a river and has one person on its back."}
|
||||
{"id": 39, "q_img": "000000076394.jpg", "q_text": "Find a picture that also a living room with a sofa and blue walls, but has walls of a similar color and there is a guitar."}
|
||||
{"id": 40, "q_img": "000000577363.jpg", "q_text": "Find a picture that also a white egret, but has the same color, shows no chair and the photo is taken from a closer distance."}
|
||||
{"id": 41, "q_img": "000000057799.jpg", "q_text": "Find a picture that also a green parrot in a cage, but shows no people and the photo is zoomed in."}
|
||||
{"id": 42, "q_img": "000000018405.jpg", "q_text": "Find a picture that also a stop sign, but shows no buildings and there are trees on both sides of the road."}
|
||||
{"id": 43, "q_img": "000000049671.jpg", "q_text": "Find a picture that also a person wearing a shirt, a tie and a hat, but shows two people that are wearing a similar outfit."}
|
||||
{"id": 44, "q_img": "000000460358.jpg", "q_text": "Find a picture that also a person throwing a baseball at a child who is holding a bat, but is set in a wood and the photo is taken from a different angle."}
|
||||
{"id": 45, "q_img": "000000177458.jpg", "q_text": "Find a picture that also a greyscale close-up portrait of a person looking at the camera, but is a woman instead of a man."}
|
||||
{"id": 46, "q_img": "000000187578.jpg", "q_text": "Find a picture that also a child with a teddy bear, but is shaped like a backpack and is being worn."}
|
||||
{"id": 47, "q_img": "000000363900.jpg", "q_text": "Find a picture that also a handbag with its contents displayed next to it, but has a carpet instead of grass and the photo is shot from the same angle."}
|
||||
{"id": 48, "q_img": "000000268574.jpg", "q_text": "Find a picture that also a lot of scissors stacked messily, but are more colorful, have a different shape and they are seen from the top."}
|
||||
{"id": 49, "q_img": "000000218732.jpg", "q_text": "Find a picture that also children wearing school uniforms seen from a close distance, but is shot in greyscale from a similar distance and shows more of them."}
|
||||
{"id": 50, "q_img": "000000038542.jpg", "q_text": "Find a picture that also multiple people eating at a table with rain umbrellas, but are eating in a garden and the umbrellas are colored."}
|
||||
{"id": 51, "q_img": "000000012900.jpg", "q_text": "Find a picture that also several wine glasses next to each other, but shows a person filling the glasses and the wine is white instead of red."}
|
||||
{"id": 52, "q_img": "000000575061.jpg", "q_text": "Find a picture that also a car in a snowy street, but is shot from the same angle and the car is a red pickup."}
|
||||
{"id": 53, "q_img": "000000566434.jpg", "q_text": "Find a picture that also a rusty truck, but is viewed from a closer distance and from the side."}
|
||||
{"id": 54, "q_img": "000000343973.jpg", "q_text": "Find a picture that also a person sitting outdoors and using a laptop, but is resting on a table and the photo shows only one person in the foreground."}
|
||||
{"id": 55, "q_img": "000000312971.jpg", "q_text": "Find a picture that also three people using a computer seen from the front, but shows the same number of people sitting on a bed and bare feet in the foreground."}
|
||||
{"id": 56, "q_img": "000000128913.jpg", "q_text": "Find a picture that also several snowboards on a wall, but shows no people."}
|
||||
{"id": 57, "q_img": "000000481114.jpg", "q_text": "Find a picture that also a double bed with turned on lamps on the sides, but has plain white sheets and the room looks cheaper."}
|
||||
{"id": 58, "q_img": "000000342900.jpg", "q_text": "Find a picture that also a person riding a horse, but is surrounded by snow and the trees are bare."}
|
||||
{"id": 59, "q_img": "000000507897.jpg", "q_text": "Find a picture that also a dark horse, but has a similar color and has fog in the background."}
|
||||
{"id": 60, "q_img": "000000460017.jpg", "q_text": "Find a picture that also a plushy toy on a double bed, but has more toys and the blankets are darker."}
|
||||
{"id": 61, "q_img": "000000526851.jpg", "q_text": "Find a picture that also a man standing next to a semaphore, but has a red semaphore instead of a green one."}
|
||||
{"id": 62, "q_img": "000000080239.jpg", "q_text": "Find a picture that also a kitchen counter, but has an open laptop with the screen visible on it."}
|
||||
{"id": 63, "q_img": "000000199973.jpg", "q_text": "Find a picture that also a Dell laptop on a table, but is of the same brand and is placed on an outdoor table."}
|
||||
{"id": 64, "q_img": "000000407154.jpg", "q_text": "Find a picture that also a kid with a baseball bat, but is shot from behind and shows the catcher."}
|
||||
{"id": 65, "q_img": "000000114181.jpg", "q_text": "Find a picture that also a fridge with magnets, but has a different color and a magnet depicting a flag attached to it."}
|
||||
{"id": 66, "q_img": "000000297636.jpg", "q_text": "Find a picture that also a mixer, but is empty and there are several bottles of alcohol next to it."}
|
||||
{"id": 67, "q_img": "000000287965.jpg", "q_text": "Find a picture that also a vehicle filled with bananas, but has a car instead of a truck."}
|
||||
{"id": 68, "q_img": "000000433177.jpg", "q_text": "Find a picture that also people looking at their laptop, but is set on a train instead of a bedroom."}
|
||||
{"id": 69, "q_img": "000000434831.jpg", "q_text": "Find a picture that also a close-up bowl filled with fruit, but shows only one type of fruit."}
|
||||
{"id": 70, "q_img": "000000007731.jpg", "q_text": "Find a picture that also a close-up feature phone, but is only one and is inside a display case."}
|
||||
{"id": 71, "q_img": "000000556629.jpg", "q_text": "Find a picture that also scissors, but shows more of them and they are for sale."}
|
||||
{"id": 72, "q_img": "000000426704.jpg", "q_text": "Find a picture that also a man looking at a computer screen sitting at a desk, but shows a plate with food and has a whiter background."}
|
||||
{"id": 73, "q_img": "000000430426.jpg", "q_text": "Find a picture that also a garbage truck, but is of a different color and is viewed from the back."}
|
||||
{"id": 74, "q_img": "000000425417.jpg", "q_text": "Find a picture that also sliced carrots, but are sliced in a different shapes and inside a clear glass bowl."}
|
||||
{"id": 75, "q_img": "000000393793.jpg", "q_text": "Find a picture that also a fire hydrant, but has a different color and is surrounded by flowers."}
|
||||
{"id": 76, "q_img": "000000527187.jpg", "q_text": "Find a picture that also a person holding a white hair drier, but is an older woman and there is a door in the background."}
|
||||
{"id": 77, "q_img": "000000575542.jpg", "q_text": "Find a picture that also a white refrigerator seen from the front, but has colorful magnets shaped like letters and a few photos."}
|
||||
{"id": 78, "q_img": "000000255336.jpg", "q_text": "Find a picture that also a person in a hospital bed, but is in color and shows the person using a phone."}
|
||||
{"id": 79, "q_img": "000000282212.jpg", "q_text": "Find a picture that also a tennis player serving on a grass court, but is a man, is dressed in the same color and has his feet in the air."}
|
||||
{"id": 80, "q_img": "000000502080.jpg", "q_text": "Find a picture that also a kitchen with an island, but has a suspended cookware display and the walls are not white."}
|
||||
{"id": 81, "q_img": "000000482717.jpg", "q_text": "Find a picture that also a woman wearing rubber boots, but has similar shoes and is holding an umbrella."}
|
||||
{"id": 82, "q_img": "000000503659.jpg", "q_text": "Find a picture that also a female tennis player, but is playing on clay and is hitting a backhand volley."}
|
||||
{"id": 83, "q_img": "000000092103.jpg", "q_text": "Find a picture that also a man standing and playing with the Wii, but has shorter hair, is wearing a hoodie and no other people are visible."}
|
||||
{"id": 84, "q_img": "000000538354.jpg", "q_text": "Find a picture that also a brown kitchen, but has the same color but has green tiles."}
|
||||
{"id": 85, "q_img": "000000121703.jpg", "q_text": "Find a picture that also a single slice of pizza on a plate, but has more toppings and is next to a hot dog."}
|
||||
{"id": 86, "q_img": "000000486024.jpg", "q_text": "Find a picture that also an open white fridge, but is full of bottles instead of food."}
|
||||
{"id": 87, "q_img": "000000191601.jpg", "q_text": "Find a picture that also a birthday cake seen from the top, but has the shape of two numbers and is shot from the same angle."}
|
||||
{"id": 88, "q_img": "000000256907.jpg", "q_text": "Find a picture that also an umbrellas roof, but is shot indoors and the umbrellas are more colored."}
|
||||
{"id": 89, "q_img": "000000378336.jpg", "q_text": "Find a picture that also a kitchen with appliances, but is more modern and white and grey."}
|
||||
{"id": 90, "q_img": "000000242024.jpg", "q_text": "Find a picture that also a zebra, but is in the middle of a road and there is a car in the background."}
|
||||
{"id": 91, "q_img": "000000349993.jpg", "q_text": "Find a picture that also a fridge, but is white and has a microwave oven resting on the counter next to it."}
|
||||
{"id": 92, "q_img": "000000501301.jpg", "q_text": "Find a picture that also three suitcases, but are seen from the top and are the same number."}
|
||||
{"id": 93, "q_img": "000000002505.jpg", "q_text": "Find a picture that also a young girl riding a horse seen from the front, but is inside an enclosure and a blue sky is visible in the background."}
|
||||
{"id": 94, "q_img": "000000177844.jpg", "q_text": "Find a picture that also a furniture store, but has more beds and shows no people on them."}
|
||||
{"id": 95, "q_img": "000000223021.jpg", "q_text": "Find a picture that also a girl laying on a bed, but shows only one girl and she is using a laptop."}
|
||||
{"id": 96, "q_img": "000000483524.jpg", "q_text": "Find a picture that also two horses pulling a wagon with people on it, but has white animals and palaces in the background."}
|
||||
{"id": 97, "q_img": "000000398147.jpg", "q_text": "Find a picture that also a fire hydrant painted with the design of a flag, but is painted with a different flag."}
|
||||
{"id": 98, "q_img": "000000408038.jpg", "q_text": "Find a picture that also an intersection with cars, but are seen from a higher angle and the photo shows a roundabout."}
|
||||
{"id": 99, "q_img": "000000555735.jpg", "q_text": "Find a picture that also an old truck, but is of a different color and has the hood open."}
|
||||
{"id": 100, "q_img": "000000034952.jpg", "q_text": "Find a picture that also a child wearing a baseball hat and a baseball glove in the foreground seen from the front, but is wearing a t-shirt of a different color, the photo has no people in the background and is shot from the same angle."}
|
||||
{"id": 101, "q_img": "000000427950.jpg", "q_text": "Find a picture that also a living room with a sofa and a TV, but shows the entrance door and has no animals."}
|
||||
{"id": 102, "q_img": "000000412231.jpg", "q_text": "Find a picture that also a person on a snowboard in the air, but has colorful pants and one hand over his head."}
|
||||
{"id": 103, "q_img": "000000401920.jpg", "q_text": "Find a picture that also a person biting a banana in the foreground, but is a child and is looking at the camera."}
|
||||
{"id": 104, "q_img": "000000080132.jpg", "q_text": "Find a picture that also a stove, but is taken from a closer distance and the front and has a kettle on it."}
|
||||
{"id": 105, "q_img": "000000433011.jpg", "q_text": "Find a picture that also a bench shot in greyscale, but is in the foreground and has at least a bird on it."}
|
||||
{"id": 106, "q_img": "000000159383.jpg", "q_text": "Find a picture that also a scooter with nobody on it seen from the side, but is white without a windshield and has a similar shape."}
|
||||
{"id": 107, "q_img": "000000026618.jpg", "q_text": "Find a picture that also a graffiti on the side of a vehicle, but has a truck instead of a train and is set on a city road."}
|
||||
{"id": 108, "q_img": "000000275670.jpg", "q_text": "Find a picture that also a little girl on a bed, but is playing with a tablet."}
|
||||
{"id": 109, "q_img": "000000330946.jpg", "q_text": "Find a picture that also a dog on a surfboard, but has a person standing behind it and is seen from a farther distance."}
|
||||
{"id": 110, "q_img": "000000538392.jpg", "q_text": "Find a picture that also a bus with people sitting seen from the back, but is open-top and the photo is shot from the same angle."}
|
||||
{"id": 111, "q_img": "000000533939.jpg", "q_text": "Find a picture that also a person behind a kitchen counter seen from the front, but has one man instead of two women and the photo is not taken on the set of a TV show."}
|
||||
{"id": 112, "q_img": "000000159689.jpg", "q_text": "Find a picture that also a person in public transport with luggage, but shows only one person and more luggage."}
|
||||
{"id": 113, "q_img": "000000258845.jpg", "q_text": "Find a picture that also a close-up person holding a glass of red wine in their hand, but is talking on the phone and the person is a man instead of a woman."}
|
||||
{"id": 114, "q_img": "000000492721.jpg", "q_text": "Find a picture that also a pizza on a plate, but has a child sitting at a dining table in front of it."}
|
||||
{"id": 115, "q_img": "000000084632.jpg", "q_text": "Find a picture that also a dog, but is wearing a life jacket and is near a body of water."}
|
||||
{"id": 116, "q_img": "000000572797.jpg", "q_text": "Find a picture that also a man seen from the back, but is on the shoreline, has a surfboard next to him and there is the sea in the background."}
|
||||
{"id": 117, "q_img": "000000533954.jpg", "q_text": "Find a picture that also a plane at the airport with a boarding bridge attached to it, but is of a different color and there is snow on the ground."}
|
||||
{"id": 118, "q_img": "000000355004.jpg", "q_text": "Find a picture that also a bedroom with two windows, but has lighter walls and a ceiling fan."}
|
||||
{"id": 119, "q_img": "000000079809.jpg", "q_text": "Find a picture that also a person riding a surfboard, but is a woman seen from a closer distance and she is wearing a bikini."}
|
||||
{"id": 120, "q_img": "000000345525.jpg", "q_text": "Find a picture that also a person under the covers with animals on a bed, but has fewer animals and more cats."}
|
||||
{"id": 121, "q_img": "000000117384.jpg", "q_text": "Find a picture that also a cake, but is shaped like an animal and has candles on it."}
|
||||
{"id": 122, "q_img": "000000366400.jpg", "q_text": "Find a picture that also a person playing beach volley, but is shot in greyscale and shows more people playing."}
|
||||
{"id": 123, "q_img": "000000029447.jpg", "q_text": "Find a picture that also a cat looking at the camera, but is darker and is next to one pair of shoes."}
|
||||
{"id": 124, "q_img": "000000192353.jpg", "q_text": "Find a picture that also a person riding a skateboard shot in greyscale, but has no backpack and the photo shows also a single car."}
|
||||
{"id": 125, "q_img": "000000560304.jpg", "q_text": "Find a picture that also a close-up man talking on the phone, but is taken in greyscale and from the back."}
|
||||
{"id": 126, "q_img": "000000495424.jpg", "q_text": "Find a picture that also a toilet, but has two of them, the photo is taken from a different angle, shows a window and the floor is darker."}
|
||||
{"id": 127, "q_img": "000000069201.jpg", "q_text": "Find a picture that also a man talking on the phone in the foreground, but has a body of water in the background."}
|
||||
{"id": 128, "q_img": "000000034654.jpg", "q_text": "Find a picture that also an open oven in the foreground with cookware in it, but shows a pot instead of a baking tray."}
|
||||
{"id": 129, "q_img": "000000278039.jpg", "q_text": "Find a picture that also a pink umbrella held by a person, but has the same color and the photo is set next to a door."}
|
||||
{"id": 130, "q_img": "000000271109.jpg", "q_text": "Find a picture that also a wine glass filled with a light clear red beverage, but has the same color and there is no bottle next to it."}
|
||||
{"id": 131, "q_img": "000000086659.jpg", "q_text": "Find a picture that also a closed laptop on an empty table in the foreground, but has a book on it and the table is made of wood."}
|
||||
{"id": 132, "q_img": "000000185803.jpg", "q_text": "Find a picture that also a man wearing a suit checking his phone, but is younger, wears a black suit and the photo is shot from the same angle."}
|
||||
{"id": 133, "q_img": "000000260641.jpg", "q_text": "Find a picture that also a dog in the foreground reflected in a mirror in the background, but is farther from the mirror and the photo has only one of them."}
|
||||
{"id": 134, "q_img": "000000467282.jpg", "q_text": "Find a picture that also a microwave oven, but has more of them in multiple colors."}
|
||||
{"id": 135, "q_img": "000000289452.jpg", "q_text": "Find a picture that also several fridges next to each other, but are of different colors and are placed outdoors."}
|
||||
{"id": 136, "q_img": "000000476440.jpg", "q_text": "Find a picture that also a person with an umbrella indoors seen from the front, but is seen from the same angle and is wearing sunglasses."}
|
||||
{"id": 137, "q_img": "000000323897.jpg", "q_text": "Find a picture that also a yellow school bus, but has the same color and there is snow around it."}
|
||||
{"id": 138, "q_img": "000000406398.jpg", "q_text": "Find a picture that also a man making a pizza in a home kitchen, but is looking into the camera and the photo is shot from the same angle."}
|
||||
{"id": 139, "q_img": "000000013618.jpg", "q_text": "Find a picture that also a standing woman holding an umbrella seen from the front, but has more trees in the background and the umbrella is darker."}
|
||||
{"id": 140, "q_img": "000000177681.jpg", "q_text": "Find a picture that also a man skateboarding, but is wearing a helmet with the visor down and there are no traffic cones."}
|
||||
{"id": 141, "q_img": "000000124865.jpg", "q_text": "Find a picture that also an hot dog in the foreground, but is on a plate and has a bottle of ketchup in the background."}
|
||||
{"id": 142, "q_img": "000000335558.jpg", "q_text": "Find a picture that also a woman cosplaying, but has an umbrella instead of a teddy bear and the photo is shot outdoors."}
|
||||
{"id": 143, "q_img": "000000067212.jpg", "q_text": "Find a picture that also two people on top of an elephant, but is shot on a beach and has one person next to it."}
|
||||
{"id": 144, "q_img": "000000116632.jpg", "q_text": "Find a picture that also three adults cooking in a home kitchen, but has the same number of adults and there is a child."}
|
||||
{"id": 145, "q_img": "000000449128.jpg", "q_text": "Find a picture that also a birthday cake with lit candles, but shows only a person seen from the front."}
|
||||
{"id": 146, "q_img": "000000188210.jpg", "q_text": "Find a picture that also a urinal, but is taken from the front, has only one of them and there is a sign over it."}
|
||||
{"id": 147, "q_img": "000000008112.jpg", "q_text": "Find a picture that also a boy using a laptop, but has the laptop on its knees and is seen from behind."}
|
||||
{"id": 148, "q_img": "000000569120.jpg", "q_text": "Find a picture that also a keyboard and a mouse on a solid background, but are Apple products and the photo has a different background color."}
|
||||
{"id": 149, "q_img": "000000155609.jpg", "q_text": "Find a picture that also three horses on a beach, but has the same number of animals and no people and no boats in the background."}
|
||||
{"id": 150, "q_img": "000000361882.jpg", "q_text": "Find a picture that also a white private jet, but has a single car next to it."}
|
||||
{"id": 151, "q_img": "000000431955.jpg", "q_text": "Find a picture that also a food truck, but shows the sky in the background, has more of them and they are open."}
|
||||
{"id": 152, "q_img": "000000011125.jpg", "q_text": "Find a picture that also a cat on a bed next to a book, but is open with its back on top."}
|
||||
{"id": 153, "q_img": "000000490296.jpg", "q_text": "Find a picture that also a raw turkey, but is inside an open oven and the photo is zoomed in and shows no people."}
|
||||
{"id": 154, "q_img": "000000241335.jpg", "q_text": "Find a picture that also a pair of horses pulling a wagon, but has only two of them, the photo is taken from a different angle, has a bluer sky and there are people on the wagon."}
|
||||
{"id": 155, "q_img": "000000113301.jpg", "q_text": "Find a picture that also two traffic lights next to each other, but are both green and the photo has a blue sky background."}
|
||||
{"id": 156, "q_img": "000000260777.jpg", "q_text": "Find a picture that also a suitcase with animals in it, but shows two children instead of cats."}
|
||||
{"id": 157, "q_img": "000000191428.jpg", "q_text": "Find a picture that also a cat, but has a teddy bear next to it and has its eyes fully opened."}
|
||||
{"id": 158, "q_img": "000000439019.jpg", "q_text": "Find a picture that also a brown horse with a jockey jumping a hurdle, but has the same color, is seen from a closer distance and from the front."}
|
||||
{"id": 159, "q_img": "000000108719.jpg", "q_text": "Find a picture that also a dog inside a car, but has its head out of the car and is turned away from the camera."}
|
||||
{"id": 160, "q_img": "000000168214.jpg", "q_text": "Find a picture that also a baseball pitcher throwing the ball seen from the side, but is wearing a white shirt and the photo is taken from the same angle and shows a grandstand in the background."}
|
||||
{"id": 161, "q_img": "000000472044.jpg", "q_text": "Find a picture that also several surfboards, but has more of them and they are resting on the grass."}
|
||||
{"id": 162, "q_img": "000000319550.jpg", "q_text": "Find a picture that also a car on a beach, but has more people and there is a clearer sky."}
|
||||
{"id": 163, "q_img": "000000495618.jpg", "q_text": "Find a picture that also a person with an umbrella taken in greyscale except for the umbrella, but is shot with the same style and there are more people."}
|
||||
{"id": 164, "q_img": "000000101523.jpg", "q_text": "Find a picture that also a motorbike in the foreground seen from the side, but has a different color, is seen from the same angle and the photo has a body of water in the background."}
|
||||
{"id": 165, "q_img": "000000346504.jpg", "q_text": "Find a picture that also a skier seen from the front, but wears a blue suit, is seen from the same angle and the sky is not visible."}
|
||||
{"id": 166, "q_img": "000000504767.jpg", "q_text": "Find a picture that also a ship in front of a beach, but has more beach umbrellas."}
|
||||
{"id": 167, "q_img": "000000192686.jpg", "q_text": "Find a picture that also people sitting on a bench seen from the back in front of the sea, but is zoomed in and has more people."}
|
||||
{"id": 168, "q_img": "000000258001.jpg", "q_text": "Find a picture that also a white clock tower, but is surrounded by water and the sky is clearer."}
|
||||
{"id": 169, "q_img": "000000155281.jpg", "q_text": "Find a picture that also a person dressed as a military, but has more people dressed similarly and they are sitting at a dining table."}
|
||||
{"id": 170, "q_img": "000000482737.jpg", "q_text": "Find a picture that also a living room with a fireplace, but shows an electric guitar."}
|
||||
{"id": 171, "q_img": "000000067158.jpg", "q_text": "Find a picture that also a woman and a man next to each other, but shows one of them holding an electronic device and a darker background."}
|
||||
{"id": 172, "q_img": "000000239332.jpg", "q_text": "Find a picture that also a closed flip phone next to an object, but shows a can instead of a camera and is shot from a lower angle."}
|
||||
{"id": 173, "q_img": "000000506685.jpg", "q_text": "Find a picture that also a man cutting food seen from the front, but is looking at the camera and there is a beer in front of him."}
|
||||
{"id": 174, "q_img": "000000044000.jpg", "q_text": "Find a picture that also a bathroom with a rectangular bathtub, but has the same shape, is fancier and has a window over it."}
|
||||
{"id": 175, "q_img": "000000535205.jpg", "q_text": "Find a picture that also a woman talking on the phone, but wears the hijab and there are people in the background."}
|
||||
{"id": 176, "q_img": "000000552835.jpg", "q_text": "Find a picture that also a person taking a selfie with a camera in the mirror of a motorbike, but has no cars in the background and is shot in greyscale."}
|
||||
{"id": 177, "q_img": "000000001710.jpg", "q_text": "Find a picture that also a man and girl playing with the Wii, but are much older and both wear a blue tie-shirt."}
|
||||
{"id": 178, "q_img": "000000369289.jpg", "q_text": "Find a picture that also a close-up hotdog held by a hand, but is without napkins and has more ketchup on it."}
|
||||
{"id": 179, "q_img": "000000076169.jpg", "q_text": "Find a picture that also a phone held by a hand seen from a close distance, but is seen from the front and the phone is an iPhone."}
|
||||
{"id": 180, "q_img": "000000469021.jpg", "q_text": "Find a picture that also several vases of different color, but has price tags and the photo is shot indoors."}
|
||||
{"id": 181, "q_img": "000000197869.jpg", "q_text": "Find a picture that also a food truck, but is blue, the photo has more people and is taken during the day."}
|
||||
{"id": 182, "q_img": "000000557390.jpg", "q_text": "Find a picture that also a pair of zebras, but has a giraffe and shows some trees in the background."}
|
||||
{"id": 183, "q_img": "000000521765.jpg", "q_text": "Find a picture that also a woman sitting on a bench with trees in the background, but has no buildings in the background, shows more trees, is zoomed in and the girl is not smiling."}
|
||||
{"id": 184, "q_img": "000000485244.jpg", "q_text": "Find a picture that also a surfboard held vertically leaning on the grass, but is mostly white and is held by a person."}
|
||||
{"id": 185, "q_img": "000000237588.jpg", "q_text": "Find a picture that also multiple elephants, but are in larger numbers and are on the shore of a body of water."}
|
||||
{"id": 186, "q_img": "000000414278.jpg", "q_text": "Find a picture that also a close-up octagonal stop sign, but has the same shape and is of a different color."}
|
||||
{"id": 187, "q_img": "000000440512.jpg", "q_text": "Find a picture that also a kitchen counter with bottles on it, but has more bottles."}
|
||||
{"id": 188, "q_img": "000000224782.jpg", "q_text": "Find a picture that also a cat sitting on a laptop, but has a couch in the background."}
|
||||
{"id": 189, "q_img": "000000369337.jpg", "q_text": "Find a picture that also a black cat, but has the same color and there is also a printer."}
|
||||
{"id": 190, "q_img": "000000362550.jpg", "q_text": "Find a picture that also a person walking a dog in a park, but has a kite instead of an umbrella and the photo shows a bluer sky."}
|
||||
{"id": 191, "q_img": "000000010348.jpg", "q_text": "Find a picture that also people at a table eating pizza, but has more people and the table is round."}
|
||||
{"id": 192, "q_img": "000000049089.jpg", "q_text": "Find a picture that also a train seen from the front, but has a yellow-colored locomotive, is in a more rural area and the sky is bluer."}
|
||||
{"id": 193, "q_img": "000000101556.jpg", "q_text": "Find a picture that also a snowboarder, but is falling into the water."}
|
||||
{"id": 194, "q_img": "000000247454.jpg", "q_text": "Find a picture that also a vase with flowers inside, but is made of clear blue glass."}
|
||||
{"id": 195, "q_img": "000000529535.jpg", "q_text": "Find a picture that also an old woman using a device on her own, but is using a laptop instead of a phone."}
|
||||
{"id": 196, "q_img": "000000568986.jpg", "q_text": "Find a picture that also a woman talking on the phone, but has a cappuccino in her hand and is wearing a black jacket."}
|
||||
{"id": 197, "q_img": "000000003744.jpg", "q_text": "Find a picture that also an oven, but is zoomed in and there are an oven and cooktop burners over it instead of a stove."}
|
||||
{"id": 198, "q_img": "000000043370.jpg", "q_text": "Find a picture that also an animal looking at itself in a mirror, but has a cat instead of a dog and the mirror is smaller and of a different shape."}
|
||||
{"id": 199, "q_img": "000000495913.jpg", "q_text": "Find a picture that also a person in a wheelchair on a tennis court, but has more players which are facing the camera."}
|
||||
{"id": 200, "q_img": "000000509353.jpg", "q_text": "Find a picture that also food in a showcase, but has pizzas instead of cakes and shows a person on the right."}
|
||||
{"id": 201, "q_img": "000000247053.jpg", "q_text": "Find a picture that also a double-decker bus seen from a three-quarter view, but is green and the photo is shot from the same angle."}
|
||||
{"id": 202, "q_img": "000000416783.jpg", "q_text": "Find a picture that also a man carrying a trolley looking at the camera, but wears a pouch and the photo is shot outdoor."}
|
||||
{"id": 203, "q_img": "000000466894.jpg", "q_text": "Find a picture that also a fire truck seen as a whole, but is seen from the side, has the same color and the photo shows more sky in the background."}
|
||||
{"id": 204, "q_img": "000000340582.jpg", "q_text": "Find a picture that also a close-up animal, but has a big horn sheep instead of a horse and shows rocks in the background."}
|
||||
{"id": 205, "q_img": "000000110669.jpg", "q_text": "Find a picture that also a policeman on a motorbike, but is on crosswalk stripes and the photo is taken from the same angle."}
|
||||
{"id": 206, "q_img": "000000396060.jpg", "q_text": "Find a picture that also a traffic light at sunrise, but has several cars and the sky has a similar color."}
|
||||
{"id": 207, "q_img": "000000316886.jpg", "q_text": "Find a picture that also a plate with a sandwich seen from above, but is on a wooden table with a glass of wine."}
|
||||
{"id": 208, "q_img": "000000494037.jpg", "q_text": "Find a picture that also a person seen through a glass of wine, but shows more people."}
|
||||
{"id": 209, "q_img": "000000394689.jpg", "q_text": "Find a picture that also a parrot, but is resting on the handlebar of a bicycle."}
|
||||
{"id": 210, "q_img": "000000190493.jpg", "q_text": "Find a picture that also two men in a street on a horse, but wears a cowboy hat and both horses are brown."}
|
||||
{"id": 211, "q_img": "000000551782.jpg", "q_text": "Find a picture that also two phones seen from the top, but shows a phone charger next to them."}
|
||||
{"id": 212, "q_img": "000000450103.jpg", "q_text": "Find a picture that also two people looking at a computer screen, but are standing and the image is in color."}
|
||||
{"id": 213, "q_img": "000000145435.jpg", "q_text": "Find a picture that also a steam train with a black locomotive, but is seen from the side and a farther distance."}
|
||||
{"id": 214, "q_img": "000000247574.jpg", "q_text": "Find a picture that also a man flying a kite seen from the back, but has blonde hair, the photo is zoomed in and has only the sky in the background."}
|
||||
{"id": 215, "q_img": "000000413621.jpg", "q_text": "Find a picture that also teddy bears sitting at a dining table, but has fewer plushies."}
|
||||
{"id": 216, "q_img": "000000192848.jpg", "q_text": "Find a picture that also a pair of shoes, but is on the floor instead of a bed and there is a dog next to it."}
|
||||
{"id": 217, "q_img": "000000532567.jpg", "q_text": "Find a picture that also multiple wall-mounted urinals, but are two and are mounted on a white wall."}
|
||||
{"id": 218, "q_img": "000000194953.jpg", "q_text": "Find a picture that also a yellow school bus, but has children getting on it."}
|
||||
{"id": 219, "q_img": "000000337128.jpg", "q_text": "Find a picture that also a steam train going towards the camera, but is passing under an arch structure and is seen from the same angle."}
|
||||
{"id": 220, "q_img": "000000539443.jpg", "q_text": "Find a picture that also people sitting with luggage in front of them, but is shot from the front and has plants in the background."}
|
||||
{"id": 221, "q_img": "000000142397.jpg", "q_text": "Find a picture that also a horse-drawn carriage, but shows no buildings and is set on snow."}
|
||||
{"id": 222, "q_img": "000000415356.jpg", "q_text": "Find a picture that also a closed toilet, but is in a wooden cabin."}
|
||||
{"id": 223, "q_img": "000000183739.jpg", "q_text": "Find a picture that also a person fixing a kite on the ground, but has two people fixing it and the photo is set outdoors."}
|
||||
{"id": 224, "q_img": "000000140958.jpg", "q_text": "Find a picture that also a man with glasses dressed in a suit and tie, but has a similar outfit and is next to a blackboard."}
|
||||
{"id": 225, "q_img": "000000171602.jpg", "q_text": "Find a picture that also a man on a motocross bike, but is in the air with his hands on the saddle while performing a trick."}
|
||||
{"id": 226, "q_img": "000000096302.jpg", "q_text": "Find a picture that also a man throwing a frisbee, but is wearing a light blue shirt instead of a suit."}
|
||||
{"id": 227, "q_img": "000000560706.jpg", "q_text": "Find a picture that also a red fire hydrant, but has the same color and is pouring out water."}
|
||||
{"id": 228, "q_img": "000000296142.jpg", "q_text": "Find a picture that also some parked busses, but have a different color, there are two of them and they are one behind the other."}
|
||||
{"id": 229, "q_img": "000000572788.jpg", "q_text": "Find a picture that also a parking meter, but has a similar shape, is green and seen from a slightly farther distance."}
|
||||
{"id": 230, "q_img": "000000556084.jpg", "q_text": "Find a picture that also a toilet placed on the outside, but shows more toilets and they are seen from the front."}
|
||||
{"id": 231, "q_img": "000000266467.jpg", "q_text": "Find a picture that also playground equipment, but is taken at night and there are more trees."}
|
||||
{"id": 232, "q_img": "000000046755.jpg", "q_text": "Find a picture that also a child with a closed suitcase, but shows the person sitting on the suitcase."}
|
||||
{"id": 233, "q_img": "000000390092.jpg", "q_text": "Find a picture that also a peacock, but is next to a car and the photo shows no people."}
|
||||
{"id": 234, "q_img": "000000132884.jpg", "q_text": "Find a picture that also children around a dining table, but are all girls and the photo is shot from the side."}
|
||||
{"id": 235, "q_img": "000000429515.jpg", "q_text": "Find a picture that also two people posing for a picture in front of a truck, but shows the same number of people and two dogs."}
|
||||
{"id": 236, "q_img": "000000025270.jpg", "q_text": "Find a picture that also a person talking on a cell phone while walking a dog on a leash, but wears sunglasses and is seen from the front."}
|
||||
{"id": 237, "q_img": "000000378472.jpg", "q_text": "Find a picture that also a black cat, but has the same color and is sitting on a chair on the grass."}
|
||||
{"id": 238, "q_img": "000000482727.jpg", "q_text": "Find a picture that also a man wearing a tie seen from the front, but is shot from the same angle and the man is wearing a t-shirt instead of a shirt."}
|
||||
{"id": 239, "q_img": "000000315484.jpg", "q_text": "Find a picture that also a flying airplane, but is half green and half white."}
|
||||
{"id": 240, "q_img": "000000238923.jpg", "q_text": "Find a picture that also a black dog with an object in its mouth, but has the same color and there is a plastic bottle instead of a frisbee."}
|
||||
{"id": 241, "q_img": "000000567023.jpg", "q_text": "Find a picture that also a laptop, but is on the floor, the photo shows more cables and no animals."}
|
||||
{"id": 242, "q_img": "000000355682.jpg", "q_text": "Find a picture that also a bike next to a bench, but has a book and is zoomed in."}
|
||||
{"id": 243, "q_img": "000000512744.jpg", "q_text": "Find a picture that also a statue of an elephant, but is more colorful and is surrounded by nature."}
|
||||
{"id": 244, "q_img": "000000004374.jpg", "q_text": "Find a picture that also a person brushing his teeth taking a picture in the mirror, but is a woman and the teeth are visible."}
|
||||
{"id": 245, "q_img": "000000283110.jpg", "q_text": "Find a picture that also a teddy bear, but is on top of some books."}
|
||||
{"id": 246, "q_img": "000000417897.jpg", "q_text": "Find a picture that also a luggage carousel in an airport seen from the side, but is taken from the same perspective and has people waiting for their luggage."}
|
||||
{"id": 247, "q_img": "000000405625.jpg", "q_text": "Find a picture that also a feature phone, but is on the original packaging instead of being held by a hand."}
|
||||
{"id": 248, "q_img": "000000250322.jpg", "q_text": "Find a picture that also a white sink and a white toilet in a bathroom, but is taken from the same angle and there is a washing machine."}
|
||||
{"id": 249, "q_img": "000000023389.jpg", "q_text": "Find a picture that also a PC mouse, but is in its original packaging."}
|
||||
{"id": 250, "q_img": "000000468831.jpg", "q_text": "Find a picture that also a toilet and a sink in a bathroom, but shows a mirror instead of a window and there is a fully-tiled wall that is not white."}
|
||||
{"id": 251, "q_img": "000000304135.jpg", "q_text": "Find a picture that also two zebras, but are in the same number and in the water."}
|
||||
{"id": 252, "q_img": "000000371460.jpg", "q_text": "Find a picture that also a red train seen from the front, but has the same color, looks more modern and is stopped at a station."}
|
||||
{"id": 253, "q_img": "000000500655.jpg", "q_text": "Find a picture that also a woman hugging a giant fake bear, but is shot outdoors and the bear is of a different color."}
|
||||
{"id": 254, "q_img": "000000232731.jpg", "q_text": "Find a picture that also a cat, but has a bottle of Coca-Cola."}
|
||||
{"id": 255, "q_img": "000000227028.jpg", "q_text": "Find a picture that also a rainbow-colored rain umbrella, but is being held up and there is the sky in the background."}
|
||||
{"id": 256, "q_img": "000000132435.jpg", "q_text": "Find a picture that also fruit on a plate, but shows orange slices instead of a banana."}
|
||||
{"id": 257, "q_img": "000000358262.jpg", "q_text": "Find a picture that also a clean desk with a computer, but is mostly white and there are flowers on it."}
|
||||
{"id": 258, "q_img": "000000574048.jpg", "q_text": "Find a picture that also a clock lamppost, but is inside a mall and the photo is taken during the day."}
|
||||
{"id": 259, "q_img": "000000383486.jpg", "q_text": "Find a picture that also people sitting on the metro seen from the front, but is taken from the same perspective, has more people and nobody is wearing a mask."}
|
||||
{"id": 260, "q_img": "000000261017.jpg", "q_text": "Find a picture that also a man seen from the front sitting in a restaurant with lots of dishes in front of him, but is shot from the same angle and there are more plates."}
|
||||
{"id": 261, "q_img": "000000043426.jpg", "q_text": "Find a picture that also a wood-fired pizza oven, but is seen up close, there is a pizza cooking and no people are visible."}
|
||||
{"id": 262, "q_img": "000000158650.jpg", "q_text": "Find a picture that also multiple fruits on a tree, but is zoomed in and shows apples instead of bananas."}
|
||||
{"id": 263, "q_img": "000000114395.jpg", "q_text": "Find a picture that also several oranges, but are fewer and in a sink."}
|
||||
{"id": 264, "q_img": "000000456139.jpg", "q_text": "Find a picture that also a bathtub, but has a toilet next to it and has a rectangular shape, the photo shows a window over it and no sink."}
|
||||
{"id": 265, "q_img": "000000550464.jpg", "q_text": "Find a picture that also a clock, but has only one of them and is set on a beach."}
|
||||
{"id": 266, "q_img": "000000462415.jpg", "q_text": "Find a picture that also a bus seen from the side, but has a different color and is on a bridge."}
|
||||
{"id": 267, "q_img": "000000571096.jpg", "q_text": "Find a picture that also a red traffic light at an intersection, but is seen through a windshield covered in droplets."}
|
||||
{"id": 268, "q_img": "000000493063.jpg", "q_text": "Find a picture that also a man biting food, but is shot from the side and is set inside a stadium."}
|
||||
{"id": 269, "q_img": "000000556362.jpg", "q_text": "Find a picture that also a flock of sheep, but shows no people and no buildings and has snow instead of grass."}
|
||||
{"id": 270, "q_img": "000000086746.jpg", "q_text": "Find a picture that also a light brown suitcase, but has the same color and the photo has the person sitting down instead of standing up."}
|
||||
{"id": 271, "q_img": "000000215432.jpg", "q_text": "Find a picture that also an ancient vase in a display case with a person looking at it, but is of a different color and is surrounded by more people."}
|
||||
{"id": 272, "q_img": "000000151918.jpg", "q_text": "Find a picture that also a little girl dressed in pink holding a tennis racket, but is wearing the same color and is not facing the camera."}
|
||||
{"id": 273, "q_img": "000000356162.jpg", "q_text": "Find a picture that also a specific dog, but has no people, is zoomed in and the dog is looking at the camera."}
|
||||
{"id": 274, "q_img": "000000084707.jpg", "q_text": "Find a picture that also a bear, but shows only one of them and it is on a paved road."}
|
||||
{"id": 275, "q_img": "000000505974.jpg", "q_text": "Find a picture that also two blue cross street name signs, but has the same color and they are seen from a closer distance."}
|
||||
{"id": 276, "q_img": "000000324137.jpg", "q_text": "Find a picture that also a parked orange and black motorbike, but is the same color and some hills are visible in the background."}
|
||||
{"id": 277, "q_img": "000000405722.jpg", "q_text": "Find a picture that also luggage, but has more of them and they are in the trunk of a car."}
|
||||
{"id": 278, "q_img": "000000244321.jpg", "q_text": "Find a picture that also some pizza in the foreground, but is on a table with a red checked tablecloth."}
|
||||
{"id": 279, "q_img": "000000003318.jpg", "q_text": "Find a picture that also a close-up fire hydrant, but is yellow with a bit of snow on it and the photo has a different background."}
|
||||
{"id": 280, "q_img": "000000190552.jpg", "q_text": "Find a picture that also a man on a skateboard seen from the back, but is on the road and the photo shows a police car."}
|
||||
{"id": 281, "q_img": "000000321586.jpg", "q_text": "Find a picture that also a person in the foreground playing with a frisbee, but is a woman who is wearing sunglasses and more trees are visible in the background."}
|
||||
{"id": 282, "q_img": "000000074741.jpg", "q_text": "Find a picture that also a person using a laptop, but shows only one person, has no walls and the background has a plain color and is less every day."}
|
||||
{"id": 283, "q_img": "000000138260.jpg", "q_text": "Find a picture that also two white polar bears, but are on the cover of a book."}
|
||||
{"id": 284, "q_img": "000000349029.jpg", "q_text": "Find a picture that also a white refrigerator, but is half empty and there is a person next to it."}
|
||||
{"id": 285, "q_img": "000000384115.jpg", "q_text": "Find a picture that also a glass and a bottle of wine, but are in front of a microwave."}
|
||||
{"id": 286, "q_img": "000000026502.jpg", "q_text": "Find a picture that also a queen-size bed, but shows lights under it and the photo is darker."}
|
||||
{"id": 287, "q_img": "000000257479.jpg", "q_text": "Find a picture that also oranges, but are divided into multiple fruit crates and the photo shows a vehicle."}
|
||||
{"id": 288, "q_img": "000000142064.jpg", "q_text": "Find a picture that also a person cooking in front of a kitchen counter, but is shot from the side and shows only one person."}
|
||||
{"id": 289, "q_img": "000000158252.jpg", "q_text": "Find a picture that also people gathered in a living room, but has a bed and there are fewer people."}
|
||||
{"id": 290, "q_img": "000000233494.jpg", "q_text": "Find a picture that also a man talking on the phone while using a laptop, but is sitting and the photo is taken in color."}
|
||||
{"id": 291, "q_img": "000000132603.jpg", "q_text": "Find a picture that also an owl resting on the hand of a person, but has its head turned away from the camera."}
|
||||
{"id": 292, "q_img": "000000313262.jpg", "q_text": "Find a picture that also a rowing boat with people on it, but has more of them and there is a flag."}
|
||||
{"id": 293, "q_img": "000000442768.jpg", "q_text": "Find a picture that also a close-up small brownish bird, but is of the same color, in the same number and is placed on a chair."}
|
||||
{"id": 294, "q_img": "000000123887.jpg", "q_text": "Find a picture that also a boat on the ground, but is rustier and there is a lighthouse in the background."}
|
||||
{"id": 295, "q_img": "000000216411.jpg", "q_text": "Find a picture that also a Wii controller, but is only one and is of a different color."}
|
||||
{"id": 296, "q_img": "000000065607.jpg", "q_text": "Find a picture that also an outside clock, but is attached to the wall from its side and has a gold color."}
|
||||
{"id": 297, "q_img": "000000413173.jpg", "q_text": "Find a picture that also multiple people carrying a refrigerator, but is shot from a different angle and there are no flowers in the background."}
|
||||
{"id": 298, "q_img": "000000523149.jpg", "q_text": "Find a picture that also a person riding a motorbike, but is only one has there is lake in the background."}
|
||||
{"id": 299, "q_img": "000000139656.jpg", "q_text": "Find a picture that also a man shot in a half-length portrait from the front, but is wearing a red tie and the photo is taken from the same angle and outdoors."}
|
||||
{"id": 300, "q_img": "000000372291.jpg", "q_text": "Find a picture that also a clock tower, but is surrounded by water and has a bluer sky in the background."}
|
||||
{"id": 301, "q_img": "000000317518.jpg", "q_text": "Find a picture that also a red umbrella, but has the same color, shows a person and the photo is shot from the top."}
|
||||
{"id": 302, "q_img": "000000356539.jpg", "q_text": "Find a picture that also a teddy bear, but is playing the guitar and the photo is taken indoors."}
|
||||
{"id": 303, "q_img": "000000244696.jpg", "q_text": "Find a picture that also some sheep, but shows more of them and there is a body of water in the foreground."}
|
||||
{"id": 304, "q_img": "000000031395.jpg", "q_text": "Find a picture that also two cows on a beach, but has the same number of them and is taken from the side."}
|
||||
{"id": 305, "q_img": "000000566570.jpg", "q_text": "Find a picture that also a close-up soup, but is inside a white bowl and the photo shows a spoon in it and other food next to it."}
|
||||
{"id": 306, "q_img": "000000096410.jpg", "q_text": "Find a picture that also a person on a couch, but shows three people, two of whom are eating."}
|
||||
{"id": 307, "q_img": "000000301688.jpg", "q_text": "Find a picture that also a person holding an umbrella under the rain in greyscale, but shows also a child under the umbrella."}
|
||||
{"id": 308, "q_img": "000000173679.jpg", "q_text": "Find a picture that also a person talking on the phone, but is wearing a military uniform."}
|
||||
{"id": 309, "q_img": "000000348127.jpg", "q_text": "Find a picture that also a little girl holding an umbrella, but is pink and the girl is standing on the grass."}
|
||||
{"id": 310, "q_img": "000000251989.jpg", "q_text": "Find a picture that also a cat, but is sitting at a dining table."}
|
||||
{"id": 311, "q_img": "000000248698.jpg", "q_text": "Find a picture that also a stationary standing skier looking at the camera, but is wearing a helmet and the photo does not show the sky in the background."}
|
||||
{"id": 312, "q_img": "000000139138.jpg", "q_text": "Find a picture that also a close-up toilet, but shows a person who is throwing up in it."}
|
||||
{"id": 313, "q_img": "000000362945.jpg", "q_text": "Find a picture that also a white toilet, but has the same color and a wooden cover."}
|
||||
{"id": 314, "q_img": "000000268878.jpg", "q_text": "Find a picture that also a tennis player hitting a two-handed backhand, but shows a woman instead of a man and she is wearing a visor hat."}
|
||||
{"id": 315, "q_img": "000000410943.jpg", "q_text": "Find a picture that also a person with a laptop sitting in a living room, but is wearing a Santa hat."}
|
||||
{"id": 316, "q_img": "000000397310.jpg", "q_text": "Find a picture that also some food in the foreground in front of a keyboard, but has a keyboard instead of a laptop."}
|
||||
{"id": 317, "q_img": "000000430266.jpg", "q_text": "Find a picture that also a mattress with something laying on it, but has a dog instead of a child."}
|
||||
{"id": 318, "q_img": "000000444637.jpg", "q_text": "Find a picture that also a child cutting paper with scissors, but has only one person seen from a closer distance."}
|
||||
{"id": 319, "q_img": "000000188881.jpg", "q_text": "Find a picture that also a small bird standing on the back of a chair, but has only one bird and the photo is shot in color."}
|
||||
{"id": 320, "q_img": "000000297614.jpg", "q_text": "Find a picture that also a bear, but is white and is at the water's edge."}
|
||||
{"id": 321, "q_img": "000000208774.jpg", "q_text": "Find a picture that also two cats playing with each other, but have a different colors and they are in a bathtub."}
|
||||
{"id": 322, "q_img": "000000023885.jpg", "q_text": "Find a picture that also a black couch, but has the same color and is leaning against a wall."}
|
||||
{"id": 323, "q_img": "000000477271.jpg", "q_text": "Find a picture that also a child playing tennis, but is hitting a ball helped by a tennis instructor."}
|
||||
{"id": 324, "q_img": "000000390151.jpg", "q_text": "Find a picture that also a black umbrella on the ground outdoors, but has the same color and is a real version of the depicted object."}
|
||||
{"id": 325, "q_img": "000000400806.jpg", "q_text": "Find a picture that also a teddy bear in the foreground, but replaces one teddy bear with a baby wearing a white onesie."}
|
||||
{"id": 326, "q_img": "000000573917.jpg", "q_text": "Find a picture that also a toaster, but has a different color and is next to a window."}
|
||||
{"id": 327, "q_img": "000000048012.jpg", "q_text": "Find a picture that also a professional kitchen with cooks, but has two of them and they are making pizzas."}
|
||||
{"id": 328, "q_img": "000000495249.jpg", "q_text": "Find a picture that also a person seen from the back with a red umbrella, but has the same color and the photo is set in an alleyway."}
|
||||
{"id": 329, "q_img": "000000094107.jpg", "q_text": "Find a picture that also a dining table with candles outdoors at night, but has a darker background, fewer tables and people sitting."}
|
||||
{"id": 330, "q_img": "000000364315.jpg", "q_text": "Find a picture that also a close-up shell flip phone held by a hand, but is black and the photo is shot outdoors."}
|
||||
{"id": 331, "q_img": "000000013527.jpg", "q_text": "Find a picture that also a silver commercial kitchen with cooks cooking, but has the same color but there are two people in it."}
|
||||
{"id": 332, "q_img": "000000552763.jpg", "q_text": "Find a picture that also seagulls seen from the side, but is shot in grayscale."}
|
||||
{"id": 333, "q_img": "000000198461.jpg", "q_text": "Find a picture that also a goalkeeper kicking the ball, but is taken from the front and shows no people in the background."}
|
||||
{"id": 334, "q_img": "000000222575.jpg", "q_text": "Find a picture that also a band playing on a stage, but is shot from the front and has umbrellas in the background."}
|
||||
{"id": 335, "q_img": "000000010676.jpg", "q_text": "Find a picture that also a half-peeled banana, but is being held by a child and the photo is taken outdoors."}
|
||||
{"id": 336, "q_img": "000000104471.jpg", "q_text": "Find a picture that also a fire truck, but is blue and there are people in the background."}
|
||||
{"id": 337, "q_img": "000000231040.jpg", "q_text": "Find a picture that also a man holding multiple hot dogs, but are placed on a white plate."}
|
||||
{"id": 338, "q_img": "000000097087.jpg", "q_text": "Find a picture that also an outdoor merchandise display, but has plush toys instead of fruit."}
|
||||
{"id": 339, "q_img": "000000256717.jpg", "q_text": "Find a picture that also an object in the foreground and a TV in the background, but has flowers instead of a laptop."}
|
||||
{"id": 340, "q_img": "000000189284.jpg", "q_text": "Find a picture that also a silver professional kitchen, but is bigger and has people in it."}
|
||||
{"id": 341, "q_img": "000000352891.jpg", "q_text": "Find a picture that also laptops on a table with no visible people, but has more of them and the table is square instead of round."}
|
||||
{"id": 342, "q_img": "000000575700.jpg", "q_text": "Find a picture that also an abandoned refrigerator, but has more of them and there is a blue sky in the background."}
|
||||
{"id": 343, "q_img": "000000544174.jpg", "q_text": "Find a picture that also a horse head peeking out of a window, but has a different color and the photo is taken from a closer distance."}
|
||||
{"id": 344, "q_img": "000000105616.jpg", "q_text": "Find a picture that also a woman holding an umbrella while it is snowing, but is taken in the daytime and the woman is seen from the front."}
|
||||
{"id": 345, "q_img": "000000104598.jpg", "q_text": "Find a picture that also a person pouring wine into a wine glass, but has red wine instead of white wine and shows no other people around."}
|
||||
{"id": 346, "q_img": "000000457471.jpg", "q_text": "Find a picture that also a person riding an elephant, but is playing soccer and the sky is not visible in the background."}
|
||||
{"id": 347, "q_img": "000000476043.jpg", "q_text": "Find a picture that also a wooden surfboard, but has the same color and is held by a man."}
|
||||
{"id": 348, "q_img": "000000558853.jpg", "q_text": "Find a picture that also a group of people with skis seen from the front, but has more people, they are younger and the photo is shot in color."}
|
||||
{"id": 349, "q_img": "000000290907.jpg", "q_text": "Find a picture that also a dog, but has only one of them and it is behind a railing."}
|
||||
{"id": 350, "q_img": "000000054659.jpg", "q_text": "Find a picture that also a round mirror over a sink, but has the same shape and is only one."}
|
||||
{"id": 351, "q_img": "000000100748.jpg", "q_text": "Find a picture that also a brown horse, but is on a green field, is seen from the side and has a woman with a helmet riding it."}
|
||||
{"id": 352, "q_img": "000000156470.jpg", "q_text": "Find a picture that also a man riding a motocross bike, but is of the same type and there is snow instead of dirt."}
|
||||
{"id": 353, "q_img": "000000427112.jpg", "q_text": "Find a picture that also a cat next to an object, but has a bowl of fruit instead of a clock."}
|
||||
{"id": 354, "q_img": "000000121550.jpg", "q_text": "Find a picture that also a birthday cake, but is blue instead of brown and has Smarties."}
|
||||
{"id": 355, "q_img": "000000068362.jpg", "q_text": "Find a picture that also multiple scissors, but are of a single color and are held in one hand."}
|
||||
{"id": 356, "q_img": "000000265678.jpg", "q_text": "Find a picture that also a cat under an object, but is set outdoors and has a bench instead of an umbrella."}
|
||||
{"id": 357, "q_img": "000000266113.jpg", "q_text": "Find a picture that also a cat in front of a TV seen from the back, but has a bird instead of a soccer match."}
|
||||
{"id": 358, "q_img": "000000123504.jpg", "q_text": "Find a picture that also a pizza oven, but is grey and there is no person."}
|
||||
{"id": 359, "q_img": "000000253559.jpg", "q_text": "Find a picture that also small boats on a beach, but shows one cow in the foreground."}
|
||||
{"id": 360, "q_img": "000000494789.jpg", "q_text": "Find a picture that also a white toilet, but has the same color, is only one and is on the grass."}
|
||||
{"id": 361, "q_img": "000000193278.jpg", "q_text": "Find a picture that also a teddy bear, but is wearing glasses and is sitting in front of a computer."}
|
||||
{"id": 362, "q_img": "000000048603.jpg", "q_text": "Find a picture that also a bag in the back of a parked vehicle, but shows a motorbike instead of a bike."}
|
||||
{"id": 363, "q_img": "000000231947.jpg", "q_text": "Find a picture that also a bear, but is white and has some food in its mouth."}
|
||||
{"id": 364, "q_img": "000000026693.jpg", "q_text": "Find a picture that also a foreground smartly dressed man wearing glasses, but is wearing a hat and the photo has a plain background."}
|
||||
{"id": 365, "q_img": "000000074021.jpg", "q_text": "Find a picture that also a phone, but is seen from a closer distance and is next to a watch."}
|
||||
{"id": 366, "q_img": "000000215423.jpg", "q_text": "Find a picture that also a grey parrot, but has the same color, is eating and is seen from a closer distance."}
|
||||
{"id": 367, "q_img": "000000051252.jpg", "q_text": "Find a picture that also a person water skiing, but has two people and the sky is not visible."}
|
||||
{"id": 368, "q_img": "000000221604.jpg", "q_text": "Find a picture that also a man wearing a shirt and a tie, but is next to a window, has no jacket and the photo is zoomed in and in color."}
|
||||
{"id": 369, "q_img": "000000542742.jpg", "q_text": "Find a picture that also a single book on a bed, but is closed and the photo is zoomed in."}
|
||||
{"id": 370, "q_img": "000000464933.jpg", "q_text": "Find a picture that also a living room, but has a TV and a big red Indian carpet."}
|
||||
{"id": 371, "q_img": "000000204267.jpg", "q_text": "Find a picture that also a person on a motorbike with a gas station in the background, but has a similar background and the photo is taken from a farther distance."}
|
||||
{"id": 372, "q_img": "000000367855.jpg", "q_text": "Find a picture that also silhouettes of multiple birds, but are on a tree and no buildings are visible."}
|
||||
{"id": 373, "q_img": "000000075991.jpg", "q_text": "Find a picture that also several animals shot with a 3D effect, but shows a different animal and the image has the same effect."}
|
||||
{"id": 374, "q_img": "000000002705.jpg", "q_text": "Find a picture that also a person sitting at a desk with a laptop seen from above, but shows one notebook on the desk."}
|
||||
{"id": 375, "q_img": "000000460066.jpg", "q_text": "Find a picture that also a person playing an instrument in a house, but has more players and shows no other people."}
|
||||
{"id": 376, "q_img": "000000194104.jpg", "q_text": "Find a picture that also a person with multiple suitcases looking at the camera, but shows more people smiling at the camera."}
|
||||
{"id": 377, "q_img": "000000179337.jpg", "q_text": "Find a picture that also a person with an object on their head, but is a man and has bananas instead of a pot."}
|
||||
{"id": 378, "q_img": "000000452266.jpg", "q_text": "Find a picture that also a clamshell phone, but is on a beach and is open."}
|
||||
{"id": 379, "q_img": "000000271281.jpg", "q_text": "Find a picture that also a man taking a picture, but is taken from a farther distance and there is the sea and no buildings in the background."}
|
||||
{"id": 380, "q_img": "000000025671.jpg", "q_text": "Find a picture that also a person lying on a surfboard seen from the front, but is shot from the same angle and has only one person."}
|
||||
{"id": 381, "q_img": "000000120788.jpg", "q_text": "Find a picture that also a person holding a surfboard under his arm, but is shot at night and is taken from behind."}
|
||||
{"id": 382, "q_img": "000000123988.jpg", "q_text": "Find a picture that also people eating outdoors at dining tables with umbrellas, but is taken from a closer distance, has a mural in the background and the umbrellas have different colors."}
|
||||
{"id": 383, "q_img": "000000056056.jpg", "q_text": "Find a picture that also a stop sign, but is taken from inside a vehicle."}
|
||||
{"id": 384, "q_img": "000000111615.jpg", "q_text": "Find a picture that also an old smartphone, but is shown to the camera by a person visible in the background."}
|
||||
{"id": 385, "q_img": "000000509762.jpg", "q_text": "Find a picture that also a person using a laptop, but is taken indoors and there are three people sitting on one sofa."}
|
||||
{"id": 386, "q_img": "000000243286.jpg", "q_text": "Find a picture that also an orange, but shows only one and is on a tree."}
|
||||
{"id": 387, "q_img": "000000061987.jpg", "q_text": "Find a picture that also a parked car with a surfboard on it, but is of a different color, the photo is shot from the back and during the day."}
|
||||
{"id": 388, "q_img": "000000373814.jpg", "q_text": "Find a picture that also two laptops side by side on a table, but has a solid-colored wall in the background and the photo shows an additional monitor."}
|
||||
{"id": 389, "q_img": "000000479843.jpg", "q_text": "Find a picture that also a white microwave oven, but has the same color, there is one clock above it and no laptops are visible."}
|
||||
{"id": 390, "q_img": "000000029107.jpg", "q_text": "Find a picture that also a work meeting around a rectangular table, but has no bottles of wine and there are a lot of laptops on the table."}
|
||||
{"id": 391, "q_img": "000000334969.jpg", "q_text": "Find a picture that also players playing American football, but is zoomed in, has more players dressed in a different color, shows two referees and no people in the background."}
|
||||
{"id": 392, "q_img": "000000402606.jpg", "q_text": "Find a picture that also a person carrying a suitcase, but is a younger boy in an airport."}
|
||||
{"id": 393, "q_img": "000000166584.jpg", "q_text": "Find a picture that also a cat inside an appliance, but is inside a fridge instead of inside an oven."}
|
||||
{"id": 394, "q_img": "000000371124.jpg", "q_text": "Find a picture that also two dogs in a kitchen, but is shot from a different angle and there is also a person."}
|
||||
{"id": 395, "q_img": "000000570604.jpg", "q_text": "Find a picture that also a bottle with a wine glass next to it in the foreground, but has the glass half-filled with red wine instead of water."}
|
||||
{"id": 396, "q_img": "000000155446.jpg", "q_text": "Find a picture that also a person playing the Wii, but shows no TV and there are two people with Christmas decorations in the background."}
|
||||
{"id": 397, "q_img": "000000036621.jpg", "q_text": "Find a picture that also an animal-shaped cake, but has a teddy-bear shape."}
|
||||
{"id": 398, "q_img": "000000171098.jpg", "q_text": "Find a picture that also a sailboat in the sea, but has more of them with the same color."}
|
||||
{"id": 399, "q_img": "000000338386.jpg", "q_text": "Find a picture that also a stop sign, but has two arrows on it and the photo has a different background."}
|
||||
{"id": 400, "q_img": "000000147028.jpg", "q_text": "Find a picture that also a person taking a picture of another person, but is seen from the back, shows a smartphone instead of a camera and the screen is in the foreground."}
|
||||
{"id": 401, "q_img": "000000136710.jpg", "q_text": "Find a picture that also a toilet and a bathtub with curtains, but has the curtains closed."}
|
||||
{"id": 402, "q_img": "000000508439.jpg", "q_text": "Find a picture that also a single person using a laptop, but is wearing no headphones and a kitchen is visible in the background."}
|
||||
{"id": 403, "q_img": "000000235559.jpg", "q_text": "Find a picture that also a man looking at the camera cutting his food in a restaurant, but is wearing glasses, is eating pizza instead of a steak and there are more people in the background."}
|
||||
{"id": 404, "q_img": "000000388038.jpg", "q_text": "Find a picture that also a dog with a frisbee in its mouth, but has a white frisbee and is set on sand."}
|
||||
{"id": 405, "q_img": "000000441490.jpg", "q_text": "Find a picture that also a parked blue motorbike, but has the same color, the photo is taken from a different angle and is set inside a repair shop."}
|
||||
{"id": 406, "q_img": "000000199439.jpg", "q_text": "Find a picture that also a bike resting on the ground, but has only one of them and has a paved surface instead of grass."}
|
||||
{"id": 407, "q_img": "000000176449.jpg", "q_text": "Find a picture that also girls playing soccer, but has fog in the background and shows no buildings."}
|
||||
{"id": 408, "q_img": "000000009135.jpg", "q_text": "Find a picture that also a red traffic light seen from the front in the foreground, but has the same color, has snow on it and the background is snowy."}
|
||||
{"id": 409, "q_img": "000000348652.jpg", "q_text": "Find a picture that also a bowl with food inside next to a keyboard, but is of a different color."}
|
||||
{"id": 410, "q_img": "000000187152.jpg", "q_text": "Find a picture that also a clear glass vase with flowers in it, but are more colorful and there is a grey brick wall in the background."}
|
||||
{"id": 411, "q_img": "000000038595.jpg", "q_text": "Find a picture that also an office desk, but shows a person reading with their face covered."}
|
||||
{"id": 412, "q_img": "000000172868.jpg", "q_text": "Find a picture that also a man playing the saxophone, but is sitting and the photo is set outdoors."}
|
||||
{"id": 413, "q_img": "000000243295.jpg", "q_text": "Find a picture that also a naked red motorcycle viewed from the side, but is on the grass and is shot from the same angle."}
|
||||
{"id": 414, "q_img": "000000433949.jpg", "q_text": "Find a picture that also a man talking on the phone, but is laying on the grass and the photo shows three people."}
|
||||
{"id": 415, "q_img": "000000155898.jpg", "q_text": "Find a picture that also a cow seen from the side, but is seen from the same angle and is in the middle of a road."}
|
||||
{"id": 416, "q_img": "000000298504.jpg", "q_text": "Find a picture that also a man skateboarding, but is wearing black shorts and a helmet instead of a baseball cap."}
|
||||
{"id": 417, "q_img": "000000141315.jpg", "q_text": "Find a picture that also a dog in the driver's seat, but is visible through the lowered window."}
|
||||
{"id": 418, "q_img": "000000388095.jpg", "q_text": "Find a picture that also an animal wearing a Santa's hat, but has a dog instead of a cat."}
|
||||
{"id": 419, "q_img": "000000200632.jpg", "q_text": "Find a picture that also a double parking meter, but has a different color and the photo shows also a hand."}
|
||||
{"id": 420, "q_img": "000000064568.jpg", "q_text": "Find a picture that also a person sitting under an umbrella fishing in a river, but has only one person and is taken from a farther distance."}
|
||||
{"id": 421, "q_img": "000000136460.jpg", "q_text": "Find a picture that also a bowl of fruit, but instead of having apples and bananas has mandarins."}
|
||||
{"id": 422, "q_img": "000000071976.jpg", "q_text": "Find a picture that also a woman holding an umbrella, but is sitting on a bench."}
|
||||
{"id": 423, "q_img": "000000208893.jpg", "q_text": "Find a picture that also a plushy toy seen from the front, but is placed on a laptop instead of inside a book and is seen from the same angle."}
|
||||
{"id": 424, "q_img": "000000548851.jpg", "q_text": "Find a picture that also a woman holding an umbrella in the foreground looking at the camera, but is shot in greyscale and is taken outdoors."}
|
||||
{"id": 425, "q_img": "000000188203.jpg", "q_text": "Find a picture that also an empty church with rows of benches, but is taken from a lower angle."}
|
||||
{"id": 426, "q_img": "000000141266.jpg", "q_text": "Find a picture that also multiple children looking at a device used by a person, but are looking at a phone instead of a laptop."}
|
||||
{"id": 427, "q_img": "000000173977.jpg", "q_text": "Find a picture that also a person tasting a glass of red wine in the foreground, but is a woman and is looking at the camera."}
|
||||
{"id": 428, "q_img": "000000407141.jpg", "q_text": "Find a picture that also a sandwich cut in half next to a cappuccino, but is placed on a black plate on a red table."}
|
||||
{"id": 429, "q_img": "000000446263.jpg", "q_text": "Find a picture that also a man sitting using a laptop facing the camera, but is shot from the same angle and has a red sofa."}
|
||||
{"id": 430, "q_img": "000000062199.jpg", "q_text": "Find a picture that also a man seen from the front playing an electric guitar, but is taken from the same angle, is brighter and shows no speakers in the background."}
|
||||
{"id": 431, "q_img": "000000472340.jpg", "q_text": "Find a picture that also a bench seen from the front, but is more minimalistic and there is a white wall as a background."}
|
||||
{"id": 432, "q_img": "000000213910.jpg", "q_text": "Find a picture that also a person on a bench using a laptop, but has only one person and shows a bag next to them."}
|
||||
{"id": 433, "q_img": "000000229438.jpg", "q_text": "Find a picture that also a laptop with a sticker attached to the back, but has a different color and more stickers."}
|
||||
{"id": 434, "q_img": "000000481994.jpg", "q_text": "Find a picture that also a person flying a kite on a meadow, but has several skyscrapers in the background."}
|
||||
{"id": 435, "q_img": "000000047989.jpg", "q_text": "Find a picture that also a person using a banana as a pretend phone, but is a half-length portrait and the person is facing the camera."}
|
||||
{"id": 436, "q_img": "000000150068.jpg", "q_text": "Find a picture that also a close-up suitcase, but is only one and is semi-open."}
|
||||
{"id": 437, "q_img": "000000553921.jpg", "q_text": "Find a picture that also a person cutting a blonde child's hair, but is seen from the front."}
|
||||
{"id": 438, "q_img": "000000512440.jpg", "q_text": "Find a picture that also a girl under a clock shot in greyscale, but is holding a cat."}
|
||||
{"id": 439, "q_img": "000000031760.jpg", "q_text": "Find a picture that also a close-up phone held by a hand, but has a body of water in the background."}
|
||||
{"id": 440, "q_img": "000000324944.jpg", "q_text": "Find a picture that also a bathhouse on a crowded beach, but has yellow umbrellas and is shot from a lower height."}
|
||||
{"id": 441, "q_img": "000000155731.jpg", "q_text": "Find a picture that also a dining table with people around it, but has the same shape and there is a cake in the center."}
|
||||
{"id": 442, "q_img": "000000580993.jpg", "q_text": "Find a picture that also a yellow fire hydrant in the foreground in front of a wall, but has the same color and the photo shows no people in the background."}
|
||||
{"id": 443, "q_img": "000000175796.jpg", "q_text": "Find a picture that also an orange with a piece of cutlery in it, but has a fork instead of a knife and the photo is zoomed in."}
|
||||
{"id": 444, "q_img": "000000260769.jpg", "q_text": "Find a picture that also a circular table with a plate and a bird on it, but has a table of the same shape and the plate is empty."}
|
||||
{"id": 445, "q_img": "000000569007.jpg", "q_text": "Find a picture that also a yellow fire hydrant, but is laid on its side on the ground."}
|
||||
{"id": 446, "q_img": "000000165652.jpg", "q_text": "Find a picture that also a man sitting on a bench, but is shot from a different angle and the man has several tattoos."}
|
||||
{"id": 447, "q_img": "000000539823.jpg", "q_text": "Find a picture that also a person walking with a surfboard under their arms, but is a woman with blond hair and is on the shore."}
|
||||
{"id": 448, "q_img": "000000257516.jpg", "q_text": "Find a picture that also a specific microwave oven, but is in a kitchen and has a person next to it."}
|
||||
{"id": 449, "q_img": "000000475137.jpg", "q_text": "Find a picture that also a man standing while speaking on the phone in a house, but shows a kitchen in the background."}
|
||||
{"id": 450, "q_img": "000000184645.jpg", "q_text": "Find a picture that also a tennis match inside a stadium, but is taken from a farther distance and shows the whole court."}
|
||||
{"id": 451, "q_img": "000000421464.jpg", "q_text": "Find a picture that also a white toilet, but has the same color and the photo shows a wastebasket and a bathtub."}
|
||||
{"id": 452, "q_img": "000000052176.jpg", "q_text": "Find a picture that also a solitary person seen from behind holding a dark umbrella, but instead of being on a beach he is in a forest."}
|
||||
{"id": 453, "q_img": "000000246266.jpg", "q_text": "Find a picture that also two women on a bench seen from behind, but is in color and both women have long dark hair."}
|
||||
{"id": 454, "q_img": "000000149772.jpg", "q_text": "Find a picture that also a Blackberry, but is held in front of the camera by a person that is facing the camera."}
|
||||
{"id": 455, "q_img": "000000135395.jpg", "q_text": "Find a picture that also people working at a picnic table outdoors, but has only one table and is zoomed in."}
|
||||
{"id": 456, "q_img": "000000066197.jpg", "q_text": "Find a picture that also an opened close-up refrigerator, but has a cat sneaking into it."}
|
||||
{"id": 457, "q_img": "000000516971.jpg", "q_text": "Find a picture that also a person standing on a bench, but has more people on it and they are still."}
|
||||
{"id": 458, "q_img": "000000260032.jpg", "q_text": "Find a picture that also a dog in front of a person on a vehicle, but has a motorbike instead of a kayak."}
|
||||
{"id": 459, "q_img": "000000486314.jpg", "q_text": "Find a picture that also a lying man reading a book, but has a bed instead of a bench and the image is in color."}
|
||||
{"id": 460, "q_img": "000000139006.jpg", "q_text": "Find a picture that also a close-up white seagull, but instead of having a beach has a street in the background."}
|
||||
{"id": 461, "q_img": "000000138173.jpg", "q_text": "Find a picture that also an elderly couple, but is sitting on a bench while reading."}
|
||||
{"id": 462, "q_img": "000000166286.jpg", "q_text": "Find a picture that also two plush toys, but depict two bears resting on a couch."}
|
||||
{"id": 463, "q_img": "000000570911.jpg", "q_text": "Find a picture that also a dog, but is only one and is in a cage."}
|
||||
{"id": 464, "q_img": "000000509317.jpg", "q_text": "Find a picture that also a clock, but has a rooster over it."}
|
||||
{"id": 465, "q_img": "000000485410.jpg", "q_text": "Find a picture that also a hotel room with two beds and no people in it, but is messier and the photo is shot from a different angle."}
|
||||
{"id": 466, "q_img": "000000076800.jpg", "q_text": "Find a picture that also an object in the foreground with a turned-on TV in the background, but is taken from the same angle and shows a TV remote control instead of a Wii controller."}
|
||||
{"id": 467, "q_img": "000000567184.jpg", "q_text": "Find a picture that also multiple baskets of bananas, but are carried using one bike."}
|
||||
{"id": 468, "q_img": "000000422734.jpg", "q_text": "Find a picture that also three people playing the Wii, but has the same number of people and shows a steering wheel controller."}
|
||||
{"id": 469, "q_img": "000000562245.jpg", "q_text": "Find a picture that also multiple horses on a beach, but has one more horse and has no people."}
|
||||
{"id": 470, "q_img": "000000124818.jpg", "q_text": "Find a picture that also a person lying on a bench, but has more leaves around it."}
|
||||
{"id": 471, "q_img": "000000024902.jpg", "q_text": "Find a picture that also two people stand up paddleboarding, but has more of them."}
|
||||
{"id": 472, "q_img": "000000236427.jpg", "q_text": "Find a picture that also a kite shaped like an octopus, but is of a different color, the photo shows no people and has a bluer sky in the background."}
|
||||
{"id": 473, "q_img": "000000056947.jpg", "q_text": "Find a picture that also a dog next to a fire hydrant, but shows one more dog."}
|
||||
{"id": 474, "q_img": "000000218794.jpg", "q_text": "Find a picture that also a street with red traffic lights, but is taken from inside a car and the dashboard is visible."}
|
||||
{"id": 475, "q_img": "000000306804.jpg", "q_text": "Find a picture that also a person checking their phone, but is sitting on a bench and is seen from the front."}
|
||||
{"id": 476, "q_img": "000000541618.jpg", "q_text": "Find a picture that also an open box of sweets, but has a donut with chocolate glaze and no sprinkles and drizzle, has no people and the focus is on the box."}
|
||||
{"id": 477, "q_img": "000000014339.jpg", "q_text": "Find a picture that also a flock of sheep, but are marked with a red sign on the wool."}
|
||||
{"id": 478, "q_img": "000000426335.jpg", "q_text": "Find a picture that also a man wearing an apron in a kitchen, but is wearing a black t-shirt and has no hat."}
|
||||
{"id": 479, "q_img": "000000367921.jpg", "q_text": "Find a picture that also a famous monument, but is a famous arc and the photo is shot in Paris."}
|
||||
{"id": 480, "q_img": "000000016145.jpg", "q_text": "Find a picture that also a person talking on the phone, but is zoomed out and taken in a forest."}
|
||||
{"id": 481, "q_img": "000000339518.jpg", "q_text": "Find a picture that also a close-up person wearing a shirt and a tie, but is looking at the camera, smiling and wearing glasses and no jacket."}
|
||||
{"id": 482, "q_img": "000000499949.jpg", "q_text": "Find a picture that also an animal on an armchair, but has a footstool in front of it and the animal is a dog."}
|
||||
{"id": 483, "q_img": "000000310956.jpg", "q_text": "Find a picture that also bananas, but shows more of them and they are in a cardboard box."}
|
||||
{"id": 484, "q_img": "000000052586.jpg", "q_text": "Find a picture that also a close-up PC mouse, but has a cat instead of the keyboard."}
|
||||
{"id": 485, "q_img": "000000428750.jpg", "q_text": "Find a picture that also a person looking at their phone with other people around, but is shot in a subway carriage and is in greyscale."}
|
||||
{"id": 486, "q_img": "000000378500.jpg", "q_text": "Find a picture that also a toilet next to a bathtub of the same color, but has a different color."}
|
||||
{"id": 487, "q_img": "000000207247.jpg", "q_text": "Find a picture that also a person with an umbrella seen from the front under the snow, but has a brighter background and is taken from the same angle."}
|
||||
{"id": 488, "q_img": "000000461273.jpg", "q_text": "Find a picture that also a child holding a tennis racket, but is wearing no t-shirt and the photo is more static."}
|
||||
{"id": 489, "q_img": "000000265380.jpg", "q_text": "Find a picture that also a bear, but is sitting on a hammock and is darker."}
|
||||
{"id": 490, "q_img": "000000123255.jpg", "q_text": "Find a picture that also a clock tower, but is light green and a blue sky is visible in the background."}
|
||||
{"id": 491, "q_img": "000000061505.jpg", "q_text": "Find a picture that also a laptop in a kitchen, but is black and is seen from a farther distance."}
|
||||
{"id": 492, "q_img": "000000102280.jpg", "q_text": "Find a picture that also a pink donut next to a depiction of a Simpsons character, but has a donut of the same color and a different character from the same franchise."}
|
||||
{"id": 493, "q_img": "000000303331.jpg", "q_text": "Find a picture that also a girl talking on the phone, but is laying on the floor instead of standing and the photo is shot from the top."}
|
||||
{"id": 494, "q_img": "000000380925.jpg", "q_text": "Find a picture that also a close-up muffin, but has a decoration depicting an animal."}
|
||||
{"id": 495, "q_img": "000000330360.jpg", "q_text": "Find a picture that also a white and green laptop, but is seen from a different angle and the photo shows a cat and no mouse next to it."}
|
||||
{"id": 496, "q_img": "000000391592.jpg", "q_text": "Find a picture that also a brown cake, but has the same color, is whole and has three or more candles on it."}
|
||||
{"id": 497, "q_img": "000000561260.jpg", "q_text": "Find a picture that also a person on a motor-scooter seen from the back, but shows only one person with no backpack and there is a more rural background."}
|
||||
{"id": 498, "q_img": "000000360027.jpg", "q_text": "Find a picture that also a sink in the foreground, but is surrounded by colored walls and the photo is shot from the same angle."}
|
||||
{"id": 499, "q_img": "000000255377.jpg", "q_text": "Find a picture that also multiple clocks on a wall, but shows a person in the foreground and the wall is brighter."}
|
||||
{"id": 500, "q_img": "000000496184.jpg", "q_text": "Find a picture that also a double-decker bus viewed from the front left, but is green and white and the photo is shot from the same angle."}
|
||||
{"id": 501, "q_img": "000000096126.jpg", "q_text": "Find a picture that also a plain white plate with broccoli in the foreground, but has potatoes instead of lentils."}
|
||||
{"id": 502, "q_img": "000000075127.jpg", "q_text": "Find a picture that also a sofa in the background, but has a table and an archway divider in the foreground."}
|
||||
{"id": 503, "q_img": "000000499841.jpg", "q_text": "Find a picture that also a dog with a frisbee in its mouth, but has its paws in the water and has the same color."}
|
||||
{"id": 504, "q_img": "000000009533.jpg", "q_text": "Find a picture that also a small white refrigerator shot in greyscale, but is only one and in the corner of a room."}
|
||||
{"id": 505, "q_img": "000000422173.jpg", "q_text": "Find a picture that also a filled stem glass, but is held in a hand and there is snow in the background."}
|
||||
{"id": 506, "q_img": "000000213161.jpg", "q_text": "Find a picture that also a living room with a full-wall window, but has an armchair and the window spans two walls."}
|
||||
{"id": 507, "q_img": "000000134934.jpg", "q_text": "Find a picture that also a double-decker bus, but is green and is viewed from the front."}
|
||||
{"id": 508, "q_img": "000000048567.jpg", "q_text": "Find a picture that also a bed, but is taken from the front and has three pictures on the wall behind it."}
|
||||
{"id": 509, "q_img": "000000427006.jpg", "q_text": "Find a picture that also multiple bananas, but are on a man's head."}
|
||||
{"id": 510, "q_img": "000000179546.jpg", "q_text": "Find a picture that also people dressed in uniforms cutting a cake, but has one more person in the foreground and fewer people in the background."}
|
||||
{"id": 511, "q_img": "000000237733.jpg", "q_text": "Find a picture that also a pizza, but has four of them, they are smaller and no plates are visible."}
|
||||
{"id": 512, "q_img": "000000567163.jpg", "q_text": "Find a picture that also a Harley Davidson-style motorbike, but is shot in a parking lot and has an American flag pinned on the back."}
|
||||
{"id": 513, "q_img": "000000561659.jpg", "q_text": "Find a picture that also a man skateboarding on a road, but has the person wearing a helmet and kneepads instead of a hat, shows traffic cones and the photo is taken from the front and a closer distance."}
|
||||
{"id": 514, "q_img": "000000570329.jpg", "q_text": "Find a picture that also a teddy bear in the foreground, but has a hat instead of glasses."}
|
||||
{"id": 515, "q_img": "000000109261.jpg", "q_text": "Find a picture that also a toilet, but has a frog on the seat."}
|
||||
{"id": 516, "q_img": "000000250457.jpg", "q_text": "Find a picture that also a vase with flowers, but is next to a laptop on a desk."}
|
||||
{"id": 517, "q_img": "000000398298.jpg", "q_text": "Find a picture that also multiple donuts in the foreground, but are in larger numbers, more colorful, in a display case and no other types of sweets are visible."}
|
||||
{"id": 518, "q_img": "000000259325.jpg", "q_text": "Find a picture that also multiple animals grazing seen from a distance, but has cows instead of sheep, has some clouds in the background and is shot from a similar distance."}
|
||||
{"id": 519, "q_img": "000000219950.jpg", "q_text": "Find a picture that also a filled pitcher on a restaurant table next to a pizza, but has beer instead of coke."}
|
||||
{"id": 520, "q_img": "000000408821.jpg", "q_text": "Find a picture that also a person riding a motorbike in the foreground in greyscale, but has no backpack and the photo shows no sky in the background."}
|
||||
{"id": 521, "q_img": "000000453754.jpg", "q_text": "Find a picture that also the inside of a home oven, but has a chicken inside instead of two pizzas."}
|
||||
{"id": 522, "q_img": "000000335487.jpg", "q_text": "Find a picture that also a person seen from behind holding a colored umbrella, but has the sea in the background and there is one more person in the photo."}
|
||||
{"id": 523, "q_img": "000000068495.jpg", "q_text": "Find a picture that also an object next to a door seen from the front, but is shot from the same angle and has a single bike instead of TVs."}
|
||||
{"id": 524, "q_img": "000000342555.jpg", "q_text": "Find a picture that also a person riding a horse in the countryside, but shows more people and the horses are running."}
|
||||
{"id": 525, "q_img": "000000366928.jpg", "q_text": "Find a picture that also a red and white fire hydrant next to a dog, but has the same color and the dog is smaller."}
|
||||
{"id": 526, "q_img": "000000291992.jpg", "q_text": "Find a picture that also a grown man holding a big teddy bear seen from the front, but wears a tie and the teddy bear is brown."}
|
||||
{"id": 527, "q_img": "000000347000.jpg", "q_text": "Find a picture that also a man tennis player with black pants who is serving, but is taken from a different angle, has a white shirt and similar pants."}
|
||||
{"id": 528, "q_img": "000000368488.jpg", "q_text": "Find a picture that also a large number of toilets, but is shot indoors and they are arranged in columns."}
|
||||
{"id": 529, "q_img": "000000368361.jpg", "q_text": "Find a picture that also a tie, but has a different color and depicts a keyboard."}
|
||||
{"id": 530, "q_img": "000000423126.jpg", "q_text": "Find a picture that also a zebra, but shows only one of them in the foreground while it is eating grass and has some rocks in the background."}
|
||||
{"id": 531, "q_img": "000000128802.jpg", "q_text": "Find a picture that also a person in the foreground biting a donut with a glaze, but is shot from the same angle and the glaze is white."}
|
||||
{"id": 532, "q_img": "000000049409.jpg", "q_text": "Find a picture that also a slice of cake on a plate, but is yellow and is on a white plate."}
|
||||
{"id": 533, "q_img": "000000355099.jpg", "q_text": "Find a picture that also two women holding umbrellas seen from the front, but is taken in greyscale, shows microphones and has a darker background."}
|
||||
{"id": 534, "q_img": "000000072209.jpg", "q_text": "Find a picture that also a statue of scissors, but shows also paper and rock."}
|
||||
{"id": 535, "q_img": "000000207412.jpg", "q_text": "Find a picture that also a person holding an umbrella seen from behind in greyscale, but is shot in completely different scenery and there are two people under the umbrella."}
|
||||
{"id": 536, "q_img": "000000217483.jpg", "q_text": "Find a picture that also a stack of books in the foreground, but are on a bed and there is no stereo."}
|
||||
{"id": 537, "q_img": "000000020214.jpg", "q_text": "Find a picture that also a living room with a sofa, but is seen as a reflection in a mirror."}
|
||||
{"id": 538, "q_img": "000000275081.jpg", "q_text": "Find a picture that also a child wearing a tie, but is a girl and is seen from a farther distance."}
|
||||
{"id": 539, "q_img": "000000071118.jpg", "q_text": "Find a picture that also a woman talking on the phone seen from the front, but has a darker tank top and there is a tree in the background."}
|
||||
{"id": 540, "q_img": "000000432862.jpg", "q_text": "Find a picture that also a cat inside an object, but shows a vase instead of a bowl."}
|
||||
{"id": 541, "q_img": "000000328752.jpg", "q_text": "Find a picture that also a cross-country skier, but is seen from the side and wears a backpack."}
|
||||
{"id": 542, "q_img": "000000403928.jpg", "q_text": "Find a picture that also a child with an umbrella in the foreground, but has a cow in the background and is shot in black and white."}
|
||||
{"id": 543, "q_img": "000000511019.jpg", "q_text": "Find a picture that also a parking meter in front of a mural, but shows a different animal."}
|
||||
{"id": 544, "q_img": "000000113136.jpg", "q_text": "Find a picture that also a fridge, but has magnets placed in the shape of a smiley."}
|
||||
{"id": 545, "q_img": "000000271768.jpg", "q_text": "Find a picture that also a half-eaten donut, but is resting on a plate and the photo has a mug."}
|
||||
{"id": 546, "q_img": "000000283714.jpg", "q_text": "Find a picture that also multiple surfboards on a beach, but has people on them."}
|
||||
{"id": 547, "q_img": "000000266509.jpg", "q_text": "Find a picture that also a child on a skateboard in the foreground, but shows two of them and they are sitting on it instead of riding it."}
|
||||
{"id": 548, "q_img": "000000033767.jpg", "q_text": "Find a picture that also three glasses of wine in the foreground, but has some food next to them."}
|
||||
{"id": 549, "q_img": "000000273972.jpg", "q_text": "Find a picture that also a table with a laptop on it, but has a different color and is shaped like a triangle."}
|
||||
{"id": 550, "q_img": "000000325478.jpg", "q_text": "Find a picture that also a man with a guitar and a woman with a violin seen from the front, but are younger, the photo is set indoors, has two windows in the background and shows no other people."}
|
||||
{"id": 551, "q_img": "000000175829.jpg", "q_text": "Find a picture that also people sitting on an elephant, but has more people on it and no umbrellas."}
|
||||
{"id": 552, "q_img": "000000168466.jpg", "q_text": "Find a picture that also a person holding a surfboard vertically, but has more people and is taken outdoors."}
|
||||
{"id": 553, "q_img": "000000218962.jpg", "q_text": "Find a picture that also a person doing sport shot in greyscale, but is playing tennis instead of golf and the player is dressed fully in white."}
|
||||
{"id": 554, "q_img": "000000039510.jpg", "q_text": "Find a picture that also people standing and drinking in a kitchen at a party, but shows more people and a brighter background."}
|
||||
{"id": 555, "q_img": "000000411293.jpg", "q_text": "Find a picture that also people next to a big teddy bear, but shows three people and is zoomed in."}
|
||||
{"id": 556, "q_img": "000000453794.jpg", "q_text": "Find a picture that also a man surfing, but shows people in the foreground looking at him."}
|
||||
{"id": 557, "q_img": "000000537001.jpg", "q_text": "Find a picture that also a white toilet, but has a black cat with its head hidden in it."}
|
||||
{"id": 558, "q_img": "000000185057.jpg", "q_text": "Find a picture that also a person holding an umbrella in the snow, but shows only one person and is walking away from the camera."}
|
||||
{"id": 559, "q_img": "000000011098.jpg", "q_text": "Find a picture that also a child looking through a donut, but is set in a car."}
|
||||
{"id": 560, "q_img": "000000189616.jpg", "q_text": "Find a picture that also a XO laptop in the foreground, but is on a desk with a silver laptop next to it."}
|
||||
{"id": 561, "q_img": "000000239842.jpg", "q_text": "Find a picture that also a bike leaning against an empty bench, but is taken from a closer distance and there are no birds."}
|
||||
{"id": 562, "q_img": "000000023724.jpg", "q_text": "Find a picture that also a man seated talking on the phone in greyscale, but is seated on a bench."}
|
||||
{"id": 563, "q_img": "000000425352.jpg", "q_text": "Find a picture that also a pizza in the foreground, but is a single slice and there is a fast-food cup near it."}
|
||||
{"id": 564, "q_img": "000000436503.jpg", "q_text": "Find a picture that also a person riding a horse, but is crossing a pool of dirty water and the photo is taken from the same angle."}
|
||||
{"id": 565, "q_img": "000000182506.jpg", "q_text": "Find a picture that also a desk with multiple computers in front of a window, but shows a city skyline visible through the window."}
|
||||
{"id": 566, "q_img": "000000517529.jpg", "q_text": "Find a picture that also a woman talking on the phone, but shows no other people, is set outdoors and there is a lake in the background."}
|
||||
{"id": 567, "q_img": "000000037216.jpg", "q_text": "Find a picture that also a cat next to a pair of shoes, but has boots instead of ballet flat shoes."}
|
||||
{"id": 568, "q_img": "000000122501.jpg", "q_text": "Find a picture that also a girl wearing a cosplay, but is holding a larger teddy bear with both arms."}
|
||||
{"id": 569, "q_img": "000000355529.jpg", "q_text": "Find a picture that also a Pokemon plush toy, but shows a different character from the same franchise."}
|
||||
{"id": 570, "q_img": "000000375233.jpg", "q_text": "Find a picture that also a black bear, but has the same color, is clinging to a tree and is surrounded by leaves."}
|
||||
{"id": 571, "q_img": "000000326343.jpg", "q_text": "Find a picture that also a truck carrying goods, but has logs instead of bananas."}
|
||||
{"id": 572, "q_img": "000000080226.jpg", "q_text": "Find a picture that also a pink donut with sprinkles, but has the same color and is about to be eaten by a person."}
|
||||
{"id": 573, "q_img": "000000490170.jpg", "q_text": "Find a picture that also a close-up woman talking on the phone, but is wearing a t-shirt and multiple trees are visible in the background."}
|
||||
{"id": 574, "q_img": "000000281364.jpg", "q_text": "Find a picture that also a person wearing zombie makeup, but has four people and is taken indoors."}
|
||||
{"id": 575, "q_img": "000000299803.jpg", "q_text": "Find a picture that also small boats side by side on a river, but are white and blue and has more birds in the background."}
|
||||
{"id": 576, "q_img": "000000276350.jpg", "q_text": "Find a picture that also people playing the Wii, but are playing on a projection screen instead of a TV."}
|
||||
{"id": 577, "q_img": "000000392475.jpg", "q_text": "Find a picture that also a bear, but has one more bear and a snowy background."}
|
||||
{"id": 578, "q_img": "000000357047.jpg", "q_text": "Find a picture that also a sink with a mirror occupying the whole wall, but shows two sinks and there is bigger mirror."}
|
||||
{"id": 579, "q_img": "000000334831.jpg", "q_text": "Find a picture that also a person taking a mirror selfie behind a woman who is brushing her teeth, but has a man taking the photo instead of a woman."}
|
||||
{"id": 580, "q_img": "000000519519.jpg", "q_text": "Find a picture that also tables and chairs with umbrellas, but are on the water and there are more people."}
|
||||
{"id": 581, "q_img": "000000081914.jpg", "q_text": "Find a picture that also a child in front of donuts, but shows more people and the sweets are in a box."}
|
||||
{"id": 582, "q_img": "000000451663.jpg", "q_text": "Find a picture that also a double-decker bus in a street seen from the side, but is shot in greyscale."}
|
||||
{"id": 583, "q_img": "000000317766.jpg", "q_text": "Find a picture that also a semi-closed suitcase, but shows a child instead of a teddy bear."}
|
||||
{"id": 584, "q_img": "000000107272.jpg", "q_text": "Find a picture that also a woman in a kitchen showing food to the camera, but shows an older woman."}
|
||||
{"id": 585, "q_img": "000000235972.jpg", "q_text": "Find a picture that also a white toilet, but is closed and has a cat on the seat."}
|
||||
{"id": 586, "q_img": "000000579390.jpg", "q_text": "Find a picture that also multiple people using a laptop around a table, but has more laptops and is shot in greyscale."}
|
||||
{"id": 587, "q_img": "000000257143.jpg", "q_text": "Find a picture that also some plates with pizza and other dishes, but has a swimming pool with people swimming in the background."}
|
||||
{"id": 588, "q_img": "000000102273.jpg", "q_text": "Find a picture that also a living room with a sofa, but has multiple cardboard boxes."}
|
||||
{"id": 589, "q_img": "000000531670.jpg", "q_text": "Find a picture that also a person holding the door of a refrigerator, but is half-closed, the photo has only one person and is zoomed out."}
|
||||
{"id": 590, "q_img": "000000193530.jpg", "q_text": "Find a picture that also a teacher in front of a projection screen, but is pointing to the screen and is not facing the camera."}
|
||||
{"id": 591, "q_img": "000000527646.jpg", "q_text": "Find a picture that also a filled mixer, but is closed with a hand holding it on the top."}
|
||||
{"id": 592, "q_img": "000000567672.jpg", "q_text": "Find a picture that also a cat on a bench, but has only one black cat facing the camera."}
|
||||
{"id": 593, "q_img": "000000314856.jpg", "q_text": "Find a picture that also a TV in a living room, but shows the whole room and the room has a corner window."}
|
||||
{"id": 594, "q_img": "000000308229.jpg", "q_text": "Find a picture that also a traffic light with cars in the foreground, but is brighter and there is a rainbow in the sky."}
|
||||
{"id": 595, "q_img": "000000380036.jpg", "q_text": "Find a picture that also a yellow bird on a tree, but is shot from a farther distance and the sky is visible in the background."}
|
||||
{"id": 596, "q_img": "000000559653.jpg", "q_text": "Find a picture that also elephants walking on a city road, but is shot at night and in color."}
|
||||
{"id": 597, "q_img": "000000054973.jpg", "q_text": "Find a picture that also tables with umbrellas of the same color, but have a different color and the photo has trees in the background."}
|
||||
{"id": 598, "q_img": "000000140329.jpg", "q_text": "Find a picture that also a man with a cat on his lap, but is using a laptop."}
|
||||
{"id": 599, "q_img": "000000301384.jpg", "q_text": "Find a picture that also a close-up dog wearing a hat indoors, but is seen from the front and looking at the camera."}
|
||||
{"id": 600, "q_img": "000000234172.jpg", "q_text": "Find a picture that also two pizzas, but has two cokes next to them."}
|
||||
{"id": 601, "q_img": "000000403768.jpg", "q_text": "Find a picture that also two hot dogs, but are on one plate and have ketchup on them."}
|
||||
{"id": 602, "q_img": "000000429426.jpg", "q_text": "Find a picture that also a cubic lamp clock in a park in the foreground, but has the same shape and a different color, the photo shows a blue sky in the background."}
|
||||
{"id": 603, "q_img": "000000029502.jpg", "q_text": "Find a picture that also a close-up phone, but is only one and it is charging."}
|
||||
{"id": 604, "q_img": "000000157732.jpg", "q_text": "Find a picture that also an object on an empty bench, but shows a shoe instead of a bottle."}
|
||||
{"id": 605, "q_img": "000000526148.jpg", "q_text": "Find a picture that also an elephant in the foreground, but is painting a picture with its trunk."}
|
||||
{"id": 606, "q_img": "000000132618.jpg", "q_text": "Find a picture that also a cook with a uniform in a kitchen, but has only one person in the foreground and they are making pizzas."}
|
||||
{"id": 607, "q_img": "000000284832.jpg", "q_text": "Find a picture that also two chairs and a bench in front of a house, but has red chairs instead of white ones, there is no windows and the photo is taken from the same angle."}
|
||||
{"id": 608, "q_img": "000000320274.jpg", "q_text": "Find a picture that also a clean and empty kitchen, but is brown and has a galley style."}
|
||||
{"id": 609, "q_img": "000000268912.jpg", "q_text": "Find a picture that also a dessert on a plate next to ice cream, but has donuts instead of a slice of cake."}
|
||||
{"id": 610, "q_img": "000000182386.jpg", "q_text": "Find a picture that also a clock tower, but shows a fountain and the photo is taken from a farther distance."}
|
||||
{"id": 611, "q_img": "000000395889.jpg", "q_text": "Find a picture that also a cat sitting on an object, but is of a different color and is sitting on a toilet instead of an oven."}
|
||||
{"id": 612, "q_img": "000000470157.jpg", "q_text": "Find a picture that also a child reading in a bed, but shows two people and one of them is an adult."}
|
||||
{"id": 613, "q_img": "000000013373.jpg", "q_text": "Find a picture that also a male tennis player while serving, but is dressed in white, is playing on grass and is viewed from a closer distance."}
|
||||
{"id": 614, "q_img": "000000496084.jpg", "q_text": "Find a picture that also a man on a snowboard grinding on a rail, but is on a handrail and the photo is shot in greyscale."}
|
||||
{"id": 615, "q_img": "000000491675.jpg", "q_text": "Find a picture that also a snowboarder, but has a yellowish jacket and is seen from a closer distance."}
|
||||
{"id": 616, "q_img": "000000173862.jpg", "q_text": "Find a picture that also a close-up clock, but has a flag in the background instead of a library."}
|
||||
{"id": 617, "q_img": "000000559318.jpg", "q_text": "Find a picture that also a living room with a TV, but has red walls and parquet as the floor."}
|
||||
{"id": 618, "q_img": "000000432451.jpg", "q_text": "Find a picture that also a countdown semaphore, but is red and there is a crowded street in the background."}
|
||||
{"id": 619, "q_img": "000000396881.jpg", "q_text": "Find a picture that also oranges and orange juice on a table in the foreground, but has more of them, there is no cake and there is a swimming pool in the background."}
|
||||
{"id": 620, "q_img": "000000489260.jpg", "q_text": "Find a picture that also a closed white refrigerator, but is seen from a different angle and is outdoors."}
|
||||
{"id": 621, "q_img": "000000076787.jpg", "q_text": "Find a picture that also food on cookware on cooktop burners, but shows a pizza instead of other food."}
|
||||
{"id": 622, "q_img": "000000068376.jpg", "q_text": "Find a picture that also a woman sitting at a desk seen from the front, but is writing on a sheet of paper and the photo shows one pair of glasses."}
|
||||
{"id": 623, "q_img": "000000106903.jpg", "q_text": "Find a picture that also a man barbecuing outdoors, but has hot dogs and is set on grass."}
|
||||
{"id": 624, "q_img": "000000396240.jpg", "q_text": "Find a picture that also a street with a red traffic light, but has a rainbow in the background and the traffic light has the same color."}
|
||||
{"id": 625, "q_img": "000000087362.jpg", "q_text": "Find a picture that also breakfast food next to a cup of coffee, but has pancakes on a plate instead of the bowl."}
|
||||
{"id": 626, "q_img": "000000500033.jpg", "q_text": "Find a picture that also a woman hitting a tennis forehand, but is shot in greyscale."}
|
||||
{"id": 627, "q_img": "000000067862.jpg", "q_text": "Find a picture that also a vase containing pink flowers, but is white and shaped like a gun."}
|
||||
{"id": 628, "q_img": "000000472196.jpg", "q_text": "Find a picture that also multiple donuts, but has more of them and there is also a product label."}
|
||||
{"id": 629, "q_img": "000000476825.jpg", "q_text": "Find a picture that also a child brushing his teeth, but is standing in the hallway and is seen from a farther distance."}
|
||||
{"id": 630, "q_img": "000000204629.jpg", "q_text": "Find a picture that also a person holding an umbrella, but is walking a tightrope and has their arm stretched above his head."}
|
||||
{"id": 631, "q_img": "000000304135.jpg", "q_text": "Find a picture that also two zebras, but are in the same number and in the water."}
|
||||
{"id": 632, "q_img": "000000011225.jpg", "q_text": "Find a picture that also groceries, but are resting on a kitchen counter instead of in a fridge."}
|
||||
{"id": 633, "q_img": "000000473830.jpg", "q_text": "Find a picture that also a kitchen, but has multiple stools made of wood."}
|
||||
{"id": 634, "q_img": "000000443587.jpg", "q_text": "Find a picture that also a CRT TV, but is turned on and is in front of a window letting in sunlight."}
|
||||
{"id": 635, "q_img": "000000292204.jpg", "q_text": "Find a picture that also an animal lying on a double bed, but has a brownish dog instead of a cat."}
|
||||
{"id": 636, "q_img": "000000366397.jpg", "q_text": "Find a picture that also a brown bear, but is similar, is in the water and the photo is shot from the same angle."}
|
||||
{"id": 637, "q_img": "000000439778.jpg", "q_text": "Find a picture that also a stop sign, but has more trees in the background and is set during a snowfall."}
|
||||
{"id": 638, "q_img": "000000444511.jpg", "q_text": "Find a picture that also a teddy bear, but is inside a toy box instead of on a shelf."}
|
||||
{"id": 639, "q_img": "000000290786.jpg", "q_text": "Find a picture that also skiers and snowboarders, but has three people on a chairlift viewed from below."}
|
||||
{"id": 640, "q_img": "000000549716.jpg", "q_text": "Find a picture that also two halves of a submarine sandwich on a plate, but has bread with the same shape and has fries instead of vegetables."}
|
||||
{"id": 641, "q_img": "000000293795.jpg", "q_text": "Find a picture that also a black cat, but has the same color and is inside a bowl made of clear glass."}
|
||||
{"id": 642, "q_img": "000000182404.jpg", "q_text": "Find a picture that also a tennis racket, but is lying on the ground and has the handle facing upward."}
|
||||
{"id": 643, "q_img": "000000277580.jpg", "q_text": "Find a picture that also a stop sign and a political message, but shows people holding a flag."}
|
||||
{"id": 644, "q_img": "000000328155.jpg", "q_text": "Find a picture that also a sheep, but has a cat on it, is facing the camera and has cleaner wool."}
|
||||
{"id": 645, "q_img": "000000270545.jpg", "q_text": "Find a picture that also a white toilet seen from the front, but has a more colorful background."}
|
||||
{"id": 646, "q_img": "000000399785.jpg", "q_text": "Find a picture that also a girl using a laptop, but is not on the floor and she is wearing headphones."}
|
||||
{"id": 647, "q_img": "000000170161.jpg", "q_text": "Find a picture that also outdoor bar tables with umbrellas, but is shot from a closer distance and shows no people."}
|
||||
{"id": 648, "q_img": "000000080716.jpg", "q_text": "Find a picture that also a green war-time motorcycle, but has the same color and has a man on it."}
|
||||
{"id": 649, "q_img": "000000491537.jpg", "q_text": "Find a picture that also a man walking on the beach while holding a surfboard under his arm, but is in greyscale and shows the sky in the background."}
|
||||
{"id": 650, "q_img": "000000265265.jpg", "q_text": "Find a picture that also a skater, but has a t-shirt of the same color and there are people looking at him."}
|
||||
{"id": 651, "q_img": "000000111418.jpg", "q_text": "Find a picture that also a dog jumping to catch a frisbee next to a person, but has a red frisbee in its mouth."}
|
||||
{"id": 652, "q_img": "000000074045.jpg", "q_text": "Find a picture that also a person standing playing the Wii and a person sitting, but shows two people standing and more people sitting."}
|
||||
{"id": 653, "q_img": "000000035912.jpg", "q_text": "Find a picture that also a person with a surfboard under his arm seen from behind, but is shot from the same angle and there are parked cars on the left."}
|
||||
{"id": 654, "q_img": "000000221308.jpg", "q_text": "Find a picture that also a person with a breakfast cup, but is sitting on the sofa while eating."}
|
||||
{"id": 655, "q_img": "000000415687.jpg", "q_text": "Find a picture that also a grandfather clock in a corner, but is placed next to a window."}
|
||||
{"id": 656, "q_img": "000000432240.jpg", "q_text": "Find a picture that also a person in the foreground holding a plate with hot dogs outdoors, but is older, is bringing more hot dogs and the photo is shot from the same angle."}
|
||||
{"id": 657, "q_img": "000000416947.jpg", "q_text": "Find a picture that also a bird on a dining table, but has a different color and some people are visible in the background."}
|
||||
{"id": 658, "q_img": "000000086792.jpg", "q_text": "Find a picture that also a person playing the Nintendo Wii, but is wearing a garment related to Star Wars."}
|
||||
{"id": 659, "q_img": "000000282019.jpg", "q_text": "Find a picture that also a birthday cake with multiple lit candles, but has a child seen from the front in front of it."}
|
||||
{"id": 660, "q_img": "000000276998.jpg", "q_text": "Find a picture that also a broken phone, but is held in one hand and is seen from a closer distance."}
|
||||
{"id": 661, "q_img": "000000012388.jpg", "q_text": "Find a picture that also a laptop on a table, but has a graphics tablet next to it."}
|
||||
{"id": 662, "q_img": "000000171989.jpg", "q_text": "Find a picture that also a Blackberry-style phone held in a hand, but is black and the person is not visible."}
|
||||
{"id": 663, "q_img": "000000475201.jpg", "q_text": "Find a picture that also an old man with a cowboy hat, but is fishing instead of milking, the photo shows no animals, has the sky in the background, has one fishing rod and is zoomed out."}
|
||||
{"id": 664, "q_img": "000000574439.jpg", "q_text": "Find a picture that also a double bed, but has black support and there are white walls in the background."}
|
||||
{"id": 665, "q_img": "000000139851.jpg", "q_text": "Find a picture that also a woman with a suitcase, but is sitting on it and the photo is taken outdoors."}
|
||||
{"id": 666, "q_img": "000000343594.jpg", "q_text": "Find a picture that also a clock with an animal above it, but is hanging on the wall."}
|
||||
{"id": 667, "q_img": "000000541362.jpg", "q_text": "Find a picture that also two laptops on a desk, but has a CRT monitor next to them."}
|
||||
{"id": 668, "q_img": "000000079885.jpg", "q_text": "Find a picture that also a cat inside an object, but shows a bowl of a non-solid color instead of a cup and is taken from a lower angle."}
|
||||
{"id": 669, "q_img": "000000028412.jpg", "q_text": "Find a picture that also a city bus seen from the front-right, but is yellow instead of white and the photo is shot from the same angle."}
|
||||
{"id": 670, "q_img": "000000565498.jpg", "q_text": "Find a picture that also a tidy white kitchen with a window, but has no island and has a parquet floor."}
|
||||
{"id": 671, "q_img": "000000394646.jpg", "q_text": "Find a picture that also a donut and a filled mug, but are resting on a newspaper."}
|
||||
{"id": 672, "q_img": "000000443930.jpg", "q_text": "Find a picture that also a green and white bus, but has one double-decker bus of the same color."}
|
||||
{"id": 673, "q_img": "000000571844.jpg", "q_text": "Find a picture that also people going on skateboards on a road, but shows more people and a cloudy sky in the background."}
|
||||
{"id": 674, "q_img": "000000258082.jpg", "q_text": "Find a picture that also several teddy bears, but are piled up on top of a person."}
|
||||
{"id": 675, "q_img": "000000094771.jpg", "q_text": "Find a picture that also two newlyweds, but is in greyscale with a different background."}
|
||||
{"id": 676, "q_img": "000000272201.jpg", "q_text": "Find a picture that also a white microwave, but has the same color and is on a sidewalk."}
|
||||
{"id": 677, "q_img": "000000083927.jpg", "q_text": "Find a picture that also a elephant animal, but has only one of them, is near a trainer and the photo is shot during a show."}
|
||||
{"id": 678, "q_img": "000000176706.jpg", "q_text": "Find a picture that also a man sitting with a dog, but has only two paws on the ground."}
|
||||
{"id": 679, "q_img": "000000282743.jpg", "q_text": "Find a picture that also a fridge with something inside, but has a child instead of cats."}
|
||||
{"id": 680, "q_img": "000000297636.jpg", "q_text": "Find a picture that also a mixer, but is closed and has a fish instead of ice."}
|
||||
{"id": 681, "q_img": "000000530077.jpg", "q_text": "Find a picture that also a girl holding a teddy bear, but is lying on a bed."}
|
||||
{"id": 682, "q_img": "000000027489.jpg", "q_text": "Find a picture that also a yellow school bus, but has the same color, it is seen through a rearview mirror and there is only one of them."}
|
||||
{"id": 683, "q_img": "000000577132.jpg", "q_text": "Find a picture that also a kitchen, but has two hanging lights and the fridge has a silver color."}
|
||||
{"id": 684, "q_img": "000000515690.jpg", "q_text": "Find a picture that also a pan, but shows only one of them and it is in an oven."}
|
||||
{"id": 685, "q_img": "000000163338.jpg", "q_text": "Find a picture that also a person cutting food with a knife, but is a tomato instead of a bagel."}
|
||||
{"id": 686, "q_img": "000000409818.jpg", "q_text": "Find a picture that also a person on a snowboard rail, but is taken from a closer distance, shows a snowboard instead of skis and there are no people in the background."}
|
||||
{"id": 687, "q_img": "000000346783.jpg", "q_text": "Find a picture that also a child with a skateboard held vertically in his hands, but is shot in a skatepark and has trees in the background."}
|
||||
{"id": 688, "q_img": "000000386874.jpg", "q_text": "Find a picture that also a cat in a suitcase, but is more vintage and is not blue."}
|
||||
{"id": 689, "q_img": "000000359488.jpg", "q_text": "Find a picture that also a standing person pouring wine into a glass, but is facing the camera and there are multiple bottles behind him."}
|
||||
{"id": 690, "q_img": "000000412368.jpg", "q_text": "Find a picture that also a whistling kettle on a stove burner, but has a different color and is seen from a farther distance."}
|
||||
{"id": 691, "q_img": "000000368463.jpg", "q_text": "Find a picture that also an old woman on a chair seen from the front, but has at least a child next to her and the photo has a higher quality."}
|
||||
{"id": 692, "q_img": "000000055841.jpg", "q_text": "Find a picture that also a yellow motorbike, but has the same color and a person is riding it."}
|
||||
{"id": 693, "q_img": "000000077731.jpg", "q_text": "Find a picture that also a double parking meter, but is zoomed out and there are plants next to it."}
|
||||
{"id": 694, "q_img": "000000142003.jpg", "q_text": "Find a picture that also a person using a laptop, but is on a couch and using a mouse."}
|
||||
{"id": 695, "q_img": "000000532677.jpg", "q_text": "Find a picture that also a fire hydrant, but is in the countryside, the photo is taken from a farther distance and has a hill in the background."}
|
||||
{"id": 696, "q_img": "000000442397.jpg", "q_text": "Find a picture that also a bathroom with a sink and toilet, but has a bathtub instead of a shower and there is a white window."}
|
||||
{"id": 697, "q_img": "000000268207.jpg", "q_text": "Find a picture that also a parking meter, but has a solar panel on top and there is a person next to it."}
|
||||
{"id": 698, "q_img": "000000177209.jpg", "q_text": "Find a picture that also a red fire hydrant, but has the same color, is surrounded by leaves, is seen from a farther distance and no flowers are visible."}
|
||||
{"id": 699, "q_img": "000000074190.jpg", "q_text": "Find a picture that also a close-up open fridge, but shows a person getting food from it."}
|
||||
{"id": 700, "q_img": "000000273842.jpg", "q_text": "Find a picture that also a tree-lined avenue in a city seen from an in-road angle, but shows fewer cars and a skyscraper in the background."}
|
||||
{"id": 701, "q_img": "000000126648.jpg", "q_text": "Find a picture that also a man with an umbrella, but is surrounded by darkness and no trees are visible."}
|
||||
{"id": 702, "q_img": "000000352764.jpg", "q_text": "Find a picture that also a person looking at their phone in the foreground, but is a woman with black curly hair."}
|
||||
{"id": 703, "q_img": "000000376134.jpg", "q_text": "Find a picture that also vases seen from a close distance, but has more of them and they are made of clear glass."}
|
||||
{"id": 704, "q_img": "000000280509.jpg", "q_text": "Find a picture that also a refrigerator in the foreground, but is white and has magnets only in the upper section."}
|
||||
{"id": 705, "q_img": "000000008611.jpg", "q_text": "Find a picture that also a white fridge with photos attached to it, but has the same color and there is person in front of it."}
|
||||
{"id": 706, "q_img": "000000542721.jpg", "q_text": "Find a picture that also a professional woman tennis player hitting the ball, but is hitting a forehand and she wears a black cap."}
|
||||
{"id": 707, "q_img": "000000227650.jpg", "q_text": "Find a picture that also a fire hydrant in the foreground, but has the colors of the American flag."}
|
||||
{"id": 708, "q_img": "000000330820.jpg", "q_text": "Find a picture that also a close-up vase, but has a different shape, is more colorful and is surrounded by greenery."}
|
||||
{"id": 709, "q_img": "000000456327.jpg", "q_text": "Find a picture that also some suitcases, but has more of them, they are arranged in a tower and are located in a large room."}
|
||||
{"id": 710, "q_img": "000000581800.jpg", "q_text": "Find a picture that also a truck carrying logs, but is taken from the front and shows the sky in the background."}
|
||||
{"id": 711, "q_img": "000000362204.jpg", "q_text": "Find a picture that also a skier making a turn around a red gate, but is a professional skier and the track is bordered in blue."}
|
||||
{"id": 712, "q_img": "000000045474.jpg", "q_text": "Find a picture that also an elephant during a circus show, but has only one of them and has a darker background."}
|
||||
{"id": 713, "q_img": "000000042036.jpg", "q_text": "Find a picture that also a man wearing a tie looking at the camera, but has eyeglasses and a black shirt."}
|
||||
{"id": 714, "q_img": "000000358277.jpg", "q_text": "Find a picture that also a bottle of wine next to a filled wine glass, but has only one of them and the photo shows a microwave in the background."}
|
||||
{"id": 715, "q_img": "000000456668.jpg", "q_text": "Find a picture that also a queen-size bed, but is shot from the same angle and there is a picture on the wall behind it."}
|
||||
{"id": 716, "q_img": "000000468995.jpg", "q_text": "Find a picture that also an object half hidden in a pouch, but has scissors instead of a phone and the photo has a whiter background."}
|
||||
{"id": 717, "q_img": "000000252207.jpg", "q_text": "Find a picture that also a foreground vase, but is only one, is glossier and has a white background behind it."}
|
||||
{"id": 718, "q_img": "000000436001.jpg", "q_text": "Find a picture that also two amateur tennis players playing a double, but is shot from the same angle and both players are women."}
|
||||
{"id": 719, "q_img": "000000273636.jpg", "q_text": "Find a picture that also a pizza, but is on a work desk in front of a screen and a keyboard is visible."}
|
||||
{"id": 720, "q_img": "000000543952.jpg", "q_text": "Find a picture that also a cat, but is orange and is half inside a toilet."}
|
||||
{"id": 721, "q_img": "000000428037.jpg", "q_text": "Find a picture that also a slice of cake on a plate in the foreground next to an object, but has a teapot instead of a to-go cup."}
|
||||
{"id": 722, "q_img": "000000533322.jpg", "q_text": "Find a picture that also a kitchen sink with a window behind it, but is shot in color and the sink is black."}
|
||||
{"id": 723, "q_img": "000000252026.jpg", "q_text": "Find a picture that also a cat under an object, but shows a car instead of a bench and the photo is shot in color."}
|
||||
{"id": 724, "q_img": "000000212280.jpg", "q_text": "Find a picture that also a teddy bear on a bed, but shows more teddy bears."}
|
||||
{"id": 725, "q_img": "000000375711.jpg", "q_text": "Find a picture that also a skier, but is dressed as a bear and the photo shows more people in the background."}
|
||||
{"id": 726, "q_img": "000000528464.jpg", "q_text": "Find a picture that also a red train seen from the front, but is in a station with a lot of people."}
|
||||
{"id": 727, "q_img": "000000460472.jpg", "q_text": "Find a picture that also a fruit connected to wires, but shows a different fruit and has a glass of water."}
|
||||
{"id": 728, "q_img": "000000218905.jpg", "q_text": "Find a picture that also a cat inside an object, but is white and black and there is a suitcase instead of a sink."}
|
||||
{"id": 729, "q_img": "000000156989.jpg", "q_text": "Find a picture that also a seagull, but is shot from a closer distance and has a bridge in the background."}
|
||||
{"id": 730, "q_img": "000000487991.jpg", "q_text": "Find a picture that also a man wearing a tie and a shirt, but has no jacket and the photo is shot outdoors and in greyscale."}
|
||||
{"id": 731, "q_img": "000000571533.jpg", "q_text": "Find a picture that also a single cow in the countryside, but is of a different color, is seen from the side and the photo shows the sky in the background."}
|
||||
{"id": 732, "q_img": "000000381912.jpg", "q_text": "Find a picture that also a white toilet outdoors, but has the same color and there is a person sitting on it."}
|
||||
{"id": 733, "q_img": "000000087575.jpg", "q_text": "Find a picture that also a group of skiers standing, but shows a lake in the background."}
|
||||
{"id": 734, "q_img": "000000223108.jpg", "q_text": "Find a picture that also an abandoned microwave oven, but is in a worse condition and has no microwave door."}
|
||||
{"id": 735, "q_img": "000000391569.jpg", "q_text": "Find a picture that also multiple set round tables, but have a different color and more people sitting."}
|
||||
{"id": 736, "q_img": "000000262675.jpg", "q_text": "Find a picture that also a person holding a surfboard under their arm standing in the shallow water, but wears a black wetsuit and the surfboard has a different color."}
|
||||
{"id": 737, "q_img": "000000338273.jpg", "q_text": "Find a picture that also food on a white plate in the foreground, but shows two sandwiches with baked potatoes instead of the pizza."}
|
||||
{"id": 738, "q_img": "000000216433.jpg", "q_text": "Find a picture that also a man in the foreground putting a pizza in a pizza oven, but is taken indoors and has fewer people in the background."}
|
||||
{"id": 739, "q_img": "000000296003.jpg", "q_text": "Find a picture that also a man riding a bicycle seen from the side, but has a surfboard under his arm and the image is taken from the same angle."}
|
||||
{"id": 740, "q_img": "000000038271.jpg", "q_text": "Find a picture that also a stove, but is new and wrapped in plastic wrap and is seen from a farther distance."}
|
||||
{"id": 741, "q_img": "000000062080.jpg", "q_text": "Find a picture that also a fire hydrant, but is red and is next to a stop sign."}
|
||||
{"id": 742, "q_img": "000000415371.jpg", "q_text": "Find a picture that also an orange traffic cone, but has the same color and is on a fire hydrant."}
|
||||
{"id": 743, "q_img": "000000352543.jpg", "q_text": "Find a picture that also a toilet, but is closed and has flowers."}
|
||||
{"id": 744, "q_img": "000000385106.jpg", "q_text": "Find a picture that also a white laptop on a table, but has the same color and there is a soft toy next to it."}
|
||||
{"id": 745, "q_img": "000000337360.jpg", "q_text": "Find a picture that also plated dishes on a table, but shows no people and has more plates."}
|
||||
{"id": 746, "q_img": "000000292083.jpg", "q_text": "Find a picture that also a silhouette of an animal, but has a different animal and is set on a beach."}
|
||||
{"id": 747, "q_img": "000000308435.jpg", "q_text": "Find a picture that also a bedroom, but has one couch and a wall that is not white."}
|
||||
{"id": 748, "q_img": "000000320124.jpg", "q_text": "Find a picture that also a man wearing a tie, but is sitting and has a beverage in front of him."}
|
||||
{"id": 749, "q_img": "000000003970.jpg", "q_text": "Find a picture that also an umbrella next to chairs, but shows beach chairs instead of chairs and is set next to a pool."}
|
||||
{"id": 750, "q_img": "000000529469.jpg", "q_text": "Find a picture that also thatched umbrellas with the sea in the background, but is taken from the height of the sea."}
|
||||
{"id": 751, "q_img": "000000390713.jpg", "q_text": "Find a picture that also a lunch box with food in it and cutlery seen from the top, but has two boxes and is shot from the same angle."}
|
||||
{"id": 752, "q_img": "000000253218.jpg", "q_text": "Find a picture that also a barbecue outdoors, but is seen from a higher angle and the photo shows people and a greener background."}
|
||||
{"id": 753, "q_img": "000000212356.jpg", "q_text": "Find a picture that also a zebra, but is eating grass and is next to an ostrich."}
|
||||
{"id": 754, "q_img": "000000308542.jpg", "q_text": "Find a picture that also a white vase with flowers, but depicts a head."}
|
||||
{"id": 755, "q_img": "000000221944.jpg", "q_text": "Find a picture that also a child surfing in the foreground, but has sand instead of water."}
|
||||
{"id": 756, "q_img": "000000442780.jpg", "q_text": "Find a picture that also airport luggage carts, but shows more people and a potted plant."}
|
||||
{"id": 757, "q_img": "000000059899.jpg", "q_text": "Find a picture that also a woman holding a little girl with an umbrella, but has no water in the background."}
|
||||
{"id": 758, "q_img": "000000014742.jpg", "q_text": "Find a picture that also people cutting a ribbon at an inauguration, but is of a different color and the photo shows fewer people in the foreground."}
|
||||
{"id": 759, "q_img": "000000541061.jpg", "q_text": "Find a picture that also a woman standing in front of an airplane, but is set on the grass and the woman dressed in a different color."}
|
||||
{"id": 760, "q_img": "000000490983.jpg", "q_text": "Find a picture that also a doughnut cake with a missing part on a plate, but has the same shape and the plate has a different color."}
|
||||
{"id": 761, "q_img": "000000212613.jpg", "q_text": "Find a picture that also a plate with food seen from above, but has a wooden background and the plate is white."}
|
||||
{"id": 762, "q_img": "000000031173.jpg", "q_text": "Find a picture that also a circular close-up stop sign, but has a circular shape and the photo is taken from the same angle."}
|
||||
{"id": 763, "q_img": "000000030773.jpg", "q_text": "Find a picture that also a person in the air holding a board with his hand, but shows snowboarding instead of surfing and there is a photographer in the foreground."}
|
||||
{"id": 764, "q_img": "000000457321.jpg", "q_text": "Find a picture that also a close-up man whose face is not shown who is wearing a white shirt and a tie, but is shot from the same angle and the tie is black."}
|
||||
{"id": 765, "q_img": "000000495630.jpg", "q_text": "Find a picture that also playground equipment, but has more people."}
|
||||
{"id": 766, "q_img": "000000576352.jpg", "q_text": "Find a picture that also an orange slice, but has more of them on a plate."}
|
||||
{"id": 767, "q_img": "000000375403.jpg", "q_text": "Find a picture that also a close-up hot dog, but is held by a hand and the photo is shot on a beach."}
|
||||
{"id": 768, "q_img": "000000005518.jpg", "q_text": "Find a picture that also a knife next to a piece of food on a cutting board, but has bread instead of carrots."}
|
||||
{"id": 769, "q_img": "000000375946.jpg", "q_text": "Find a picture that also multiple clocks attached to a wall, but has clocks of a non-standard shape and the wall are of a different color."}
|
||||
{"id": 770, "q_img": "000000282469.jpg", "q_text": "Find a picture that also a hot dog in the foreground held on a hand, but has a statue instead of a stadium."}
|
||||
{"id": 771, "q_img": "000000465954.jpg", "q_text": "Find a picture that also a sink and a bathtub in a bathroom, but has a greenish-colored wall and there is a window in on the left wall."}
|
||||
{"id": 772, "q_img": "000000437777.jpg", "q_text": "Find a picture that also a person facing the camera taking a picture with a flip phone outdoors, but shows only one person and a brighter background."}
|
||||
{"id": 773, "q_img": "000000271376.jpg", "q_text": "Find a picture that also multiple donuts, but are packaged in white boxes and are of a single color."}
|
||||
{"id": 774, "q_img": "000000005645.jpg", "q_text": "Find a picture that also a man wearing a tie and a bowler hat, but is shot indoors and the man is looking at the camera."}
|
||||
{"id": 775, "q_img": "000000052896.jpg", "q_text": "Find a picture that also a black couch in a living room, but shows a kitchen in the background."}
|
||||
{"id": 776, "q_img": "000000438310.jpg", "q_text": "Find a picture that also a pizza in a box in the foreground, but has two of them and has wine instead of beer."}
|
||||
{"id": 777, "q_img": "000000283015.jpg", "q_text": "Find a picture that also several pigeons, but are sitting on a semaphore."}
|
||||
{"id": 778, "q_img": "000000335689.jpg", "q_text": "Find a picture that also a carrot, but has two of them and they are on a cutting board next to a knife."}
|
||||
{"id": 779, "q_img": "000000400473.jpg", "q_text": "Find a picture that also a person talking on the phone, but is taking a sip and the photo has a more everyday background."}
|
||||
{"id": 780, "q_img": "000000428058.jpg", "q_text": "Find a picture that also a bedroom, but is shot in greyscale and shows no books."}
|
||||
{"id": 781, "q_img": "000000436398.jpg", "q_text": "Find a picture that also a person holding freshly picked vegetables, but has a knife and the vegetable is of a different type."}
|
||||
{"id": 782, "q_img": "000000013321.jpg", "q_text": "Find a picture that also an old man with a beard wearing a hat, but shows a bench instead of a chair and there is no umbrella."}
|
||||
{"id": 783, "q_img": "000000027532.jpg", "q_text": "Find a picture that also a person standing holding an umbrella seen from the back, but has a different color and shows writing on it."}
|
||||
{"id": 784, "q_img": "000000157077.jpg", "q_text": "Find a picture that also a turned-on laptop on a desk, but shows a notebook with writings on it and the photo is shot indoors."}
|
||||
{"id": 785, "q_img": "000000129413.jpg", "q_text": "Find a picture that also a laptop seen from the front with a beverage next to it, but is shot indoors and there is a person in the background."}
|
||||
{"id": 786, "q_img": "000000392749.jpg", "q_text": "Find a picture that also a white refrigerator in the foreground, but is full of food and the photo shows only the interior."}
|
||||
{"id": 787, "q_img": "000000066260.jpg", "q_text": "Find a picture that also a small electric oven, but is turned on and has some food in it."}
|
||||
{"id": 788, "q_img": "000000087988.jpg", "q_text": "Find a picture that also a kitchen countertop, but has no visible stove and there is a laptop on it."}
|
||||
{"id": 789, "q_img": "000000070656.jpg", "q_text": "Find a picture that also a white plate with food on a wooden table, but has a banana instead of the cake."}
|
||||
{"id": 790, "q_img": "000000071231.jpg", "q_text": "Find a picture that also a laptop on a desk seen from the front, but is only one and there is a keyboard under it, the photo is shot from the same angle."}
|
||||
{"id": 791, "q_img": "000000305818.jpg", "q_text": "Find a picture that also a living room with a sofa, but shows people and has a french door in the background."}
|
||||
{"id": 792, "q_img": "000000140749.jpg", "q_text": "Find a picture that also some sheep, but are drinking water and there is no bench."}
|
||||
{"id": 793, "q_img": "000000360457.jpg", "q_text": "Find a picture that also a living room with a CRT TV, but has an open window and the photo is darker."}
|
||||
{"id": 794, "q_img": "000000272178.jpg", "q_text": "Find a picture that also a big inflatable kite on the beach, but has a whale instead of an octopus."}
|
||||
{"id": 795, "q_img": "000000195881.jpg", "q_text": "Find a picture that also something inside an open suitcase, but shows an orange cat instead of a baby and the cat has the eyes closed."}
|
||||
{"id": 796, "q_img": "000000158464.jpg", "q_text": "Find a picture that also a white fridge, but has the same color and is closed and the photo has a CRT TV."}
|
||||
{"id": 797, "q_img": "000000277154.jpg", "q_text": "Find a picture that also a group of people using laptops at a table, but has more people and the table is U-shaped."}
|
||||
{"id": 798, "q_img": "000000359573.jpg", "q_text": "Find a picture that also multiple close-up donuts, but are seen from a higher perspective and have a white glaze with sprinkles."}
|
||||
{"id": 799, "q_img": "000000222734.jpg", "q_text": "Find a picture that also a person in a kitchen looking at the camera, but is a man and is doing dishes in the kitchen sink."}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,225 @@
|
||||
import os
|
||||
import faiss
|
||||
import torch
|
||||
import logging
|
||||
import datasets
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
from typing import Optional
|
||||
from dataclasses import dataclass, field
|
||||
from transformers import HfArgumentParser
|
||||
from flag_mmret import Flag_mmret
|
||||
import json
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Args:
|
||||
model_name: str = field(
|
||||
default="BAAI/BGE-VL-large",
|
||||
metadata={'help': 'Model Name'}
|
||||
)
|
||||
result_save_path: str = field(
|
||||
default="./eval/mmret_large_circo.json",
|
||||
metadata={'help': 'Where to save the results.'}
|
||||
)
|
||||
image_dir: str = field(
|
||||
default="YOUR_COCO_IMAGE_DIRECTORY",
|
||||
metadata={'help': 'Where the images located on.'}
|
||||
)
|
||||
fp16: bool = field(
|
||||
default=False,
|
||||
metadata={'help': 'Use fp16 in inference?'}
|
||||
)
|
||||
max_query_length: int = field(
|
||||
default=64,
|
||||
metadata={'help': 'Max query length.'}
|
||||
)
|
||||
max_passage_length: int = field(
|
||||
default=77,
|
||||
metadata={'help': 'Max passage length.'}
|
||||
)
|
||||
batch_size: int = field(
|
||||
default=256,
|
||||
metadata={'help': 'Inference batch size.'}
|
||||
)
|
||||
index_factory: str = field(
|
||||
default="Flat",
|
||||
metadata={'help': 'Faiss index factory.'}
|
||||
)
|
||||
k: int = field(
|
||||
default=100,
|
||||
metadata={'help': 'How many neighbors to retrieve?'}
|
||||
)
|
||||
save_embedding: bool = field(
|
||||
default=False,
|
||||
metadata={'help': 'Save embeddings in memmap at save_dir?'}
|
||||
)
|
||||
load_embedding: bool = field(
|
||||
default=False,
|
||||
metadata={'help': 'Load embeddings from save_dir?'}
|
||||
)
|
||||
save_path: str = field(
|
||||
default="embeddings.memmap",
|
||||
metadata={'help': 'Path to save embeddings.'}
|
||||
)
|
||||
|
||||
|
||||
|
||||
def index(model: Flag_mmret, corpus: datasets.Dataset, batch_size: int = 256, max_length: int=512, index_factory: str = "Flat", save_path: str = None, save_embedding: bool = False, load_embedding: bool = False):
|
||||
"""
|
||||
1. Encode the entire corpus into dense embeddings;
|
||||
2. Create faiss index;
|
||||
3. Optionally save embeddings.
|
||||
"""
|
||||
if load_embedding:
|
||||
test = model.encode("test")
|
||||
dtype = test.dtype
|
||||
dim = len(test)
|
||||
|
||||
corpus_embeddings = np.memmap(
|
||||
save_path,
|
||||
mode="r",
|
||||
dtype=dtype
|
||||
).reshape(-1, dim)
|
||||
|
||||
else:
|
||||
|
||||
corpus_embeddings = model.encode_corpus(corpus, batch_size=batch_size, max_length=max_length, corpus_type='image')
|
||||
|
||||
dim = corpus_embeddings.shape[-1]
|
||||
|
||||
if save_embedding:
|
||||
logger.info(f"saving embeddings at {save_path}...")
|
||||
memmap = np.memmap(
|
||||
save_path,
|
||||
shape=corpus_embeddings.shape,
|
||||
mode="w+",
|
||||
dtype=corpus_embeddings.dtype
|
||||
)
|
||||
|
||||
length = corpus_embeddings.shape[0]
|
||||
# add in batch
|
||||
save_batch_size = 10000
|
||||
if length > save_batch_size:
|
||||
for i in tqdm(range(0, length, save_batch_size), leave=False, desc="Saving Embeddings"):
|
||||
j = min(i + save_batch_size, length)
|
||||
memmap[i: j] = corpus_embeddings[i: j]
|
||||
else:
|
||||
memmap[:] = corpus_embeddings
|
||||
|
||||
# create faiss index
|
||||
faiss_index = faiss.index_factory(dim, index_factory, faiss.METRIC_INNER_PRODUCT)
|
||||
|
||||
|
||||
if model.device == torch.device("cuda"):
|
||||
# co = faiss.GpuClonerOptions()
|
||||
co = faiss.GpuMultipleClonerOptions()
|
||||
co.useFloat16 = True
|
||||
# faiss_index = faiss.index_cpu_to_gpu(faiss.StandardGpuResources(), 0, faiss_index, co)
|
||||
faiss_index = faiss.index_cpu_to_all_gpus(faiss_index, co)
|
||||
|
||||
# NOTE: faiss only accepts float32
|
||||
logger.info("Adding embeddings...")
|
||||
corpus_embeddings = corpus_embeddings.astype(np.float32)
|
||||
faiss_index.train(corpus_embeddings)
|
||||
faiss_index.add(corpus_embeddings)
|
||||
|
||||
|
||||
|
||||
return faiss_index
|
||||
|
||||
|
||||
def search(model: Flag_mmret, queries: datasets, faiss_index: faiss.Index, k:int = 100, batch_size: int = 256, max_length: int=512):
|
||||
"""
|
||||
1. Encode queries into dense embeddings;
|
||||
2. Search through faiss index
|
||||
"""
|
||||
query_embeddings = model.encode_queries([queries["q_text"], queries["q_img"]],
|
||||
batch_size=batch_size,
|
||||
max_length=max_length,
|
||||
query_type='mm_it')
|
||||
|
||||
|
||||
query_size = len(query_embeddings)
|
||||
|
||||
all_scores = []
|
||||
all_indices = []
|
||||
|
||||
for i in tqdm(range(0, query_size, batch_size), desc="Searching"):
|
||||
j = min(i + batch_size, query_size)
|
||||
query_embedding = query_embeddings[i: j]
|
||||
score, indice = faiss_index.search(query_embedding.astype(np.float32), k=k)
|
||||
all_scores.append(score)
|
||||
all_indices.append(indice)
|
||||
|
||||
all_scores = np.concatenate(all_scores, axis=0)
|
||||
all_indices = np.concatenate(all_indices, axis=0)
|
||||
return all_scores, all_indices
|
||||
|
||||
|
||||
def main():
|
||||
parser = HfArgumentParser([Args])
|
||||
args: Args = parser.parse_args_into_dataclasses()[0]
|
||||
|
||||
print(f"Results will be saved in {args.result_save_path}")
|
||||
eval_data = datasets.load_dataset('json', data_files="./eval/data/circo_query.jsonl", split='train')
|
||||
image_corpus_test = datasets.load_dataset('json', data_files="./eval/data/circo_corpus.jsonl", split='train')
|
||||
|
||||
model = Flag_mmret(model_name=args.model_name,
|
||||
normlized = True,
|
||||
image_dir=args.image_dir,
|
||||
use_fp16=False,
|
||||
)
|
||||
|
||||
|
||||
faiss_index = index(
|
||||
model=model,
|
||||
corpus=image_corpus_test,
|
||||
batch_size=args.batch_size,
|
||||
max_length=args.max_passage_length,
|
||||
index_factory=args.index_factory,
|
||||
save_path=args.save_path,
|
||||
save_embedding=args.save_embedding,
|
||||
load_embedding=args.load_embedding
|
||||
)
|
||||
|
||||
scores, indices = search(
|
||||
model=model,
|
||||
queries=eval_data,
|
||||
faiss_index=faiss_index,
|
||||
k=args.k,
|
||||
batch_size=args.batch_size,
|
||||
max_length=args.max_query_length
|
||||
)
|
||||
|
||||
|
||||
retrieval_results = []
|
||||
for indice in indices:
|
||||
# filter invalid indices
|
||||
indice = indice[indice != -1].tolist()
|
||||
retrieval_results.append(image_corpus_test[indice]["content"])
|
||||
|
||||
########## results in test corpus #########
|
||||
q_images = eval_data["q_img"]
|
||||
|
||||
q_ids = []
|
||||
for _img in q_images:
|
||||
_id = os.path.basename(_img)
|
||||
_id = os.path.splitext(_id)[0]
|
||||
q_ids.append(_id)
|
||||
|
||||
pairids = eval_data["id"]
|
||||
results = {}
|
||||
for pairid, re_results, q_img in zip(pairids, retrieval_results, q_images):
|
||||
id = str(pairid)
|
||||
top_50_results = re_results[0:50]
|
||||
results[id] = top_50_results
|
||||
|
||||
with open(args.result_save_path, "w") as f:
|
||||
json.dump(results, f)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,342 @@
|
||||
import os
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = "0"
|
||||
|
||||
import sys
|
||||
print(os.getcwd())
|
||||
|
||||
import faiss
|
||||
import torch
|
||||
import logging
|
||||
import datasets
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
from typing import Optional
|
||||
from dataclasses import dataclass, field
|
||||
from transformers import HfArgumentParser
|
||||
from flag_mmret import Flag_mmret
|
||||
import json
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Args:
|
||||
model_name: str = field(
|
||||
default="BAAI/BGE-VL-large",
|
||||
metadata={'help': 'Model Name'}
|
||||
)
|
||||
image_dir: str = field(
|
||||
default="YOUR_FASHIONIQ_IMAGE_DIRECTORY",
|
||||
metadata={'help': 'Where are the images located on.'}
|
||||
)
|
||||
fp16: bool = field(
|
||||
default=False,
|
||||
metadata={'help': 'Use fp16 in inference?'}
|
||||
)
|
||||
max_query_length: int = field(
|
||||
default=64,
|
||||
metadata={'help': 'Max query length.'}
|
||||
)
|
||||
max_passage_length: int = field(
|
||||
default=77,
|
||||
metadata={'help': 'Max passage length.'}
|
||||
)
|
||||
batch_size: int = field(
|
||||
default=256,
|
||||
metadata={'help': 'Inference batch size.'}
|
||||
)
|
||||
index_factory: str = field(
|
||||
default="Flat",
|
||||
metadata={'help': 'Faiss index factory.'}
|
||||
)
|
||||
k: int = field(
|
||||
default=100,
|
||||
metadata={'help': 'How many neighbors to retrieve?'}
|
||||
)
|
||||
save_embedding: bool = field(
|
||||
default=False,
|
||||
metadata={'help': 'Save embeddings in memmap at save_dir?'}
|
||||
)
|
||||
load_embedding: bool = field(
|
||||
default=False,
|
||||
metadata={'help': 'Load embeddings from save_dir?'}
|
||||
)
|
||||
save_path: str = field(
|
||||
default="embeddings.memmap",
|
||||
metadata={'help': 'Path to save embeddings.'}
|
||||
)
|
||||
|
||||
|
||||
|
||||
def index(model: Flag_mmret, corpus: datasets.Dataset, batch_size: int = 256, max_length: int=512, index_factory: str = "Flat", save_path: str = None, save_embedding: bool = False, load_embedding: bool = False):
|
||||
"""
|
||||
1. Encode the entire corpus into dense embeddings;
|
||||
2. Create faiss index;
|
||||
3. Optionally save embeddings.
|
||||
"""
|
||||
if load_embedding:
|
||||
test = model.encode("test")
|
||||
dtype = test.dtype
|
||||
dim = len(test)
|
||||
|
||||
corpus_embeddings = np.memmap(
|
||||
save_path,
|
||||
mode="r",
|
||||
dtype=dtype
|
||||
).reshape(-1, dim)
|
||||
|
||||
else:
|
||||
|
||||
corpus_embeddings = model.encode_corpus(corpus, batch_size=batch_size, max_length=max_length, corpus_type='image')
|
||||
|
||||
dim = corpus_embeddings.shape[-1]
|
||||
|
||||
if save_embedding:
|
||||
logger.info(f"saving embeddings at {save_path}...")
|
||||
memmap = np.memmap(
|
||||
save_path,
|
||||
shape=corpus_embeddings.shape,
|
||||
mode="w+",
|
||||
dtype=corpus_embeddings.dtype
|
||||
)
|
||||
|
||||
length = corpus_embeddings.shape[0]
|
||||
# add in batch
|
||||
save_batch_size = 10000
|
||||
if length > save_batch_size:
|
||||
for i in tqdm(range(0, length, save_batch_size), leave=False, desc="Saving Embeddings"):
|
||||
j = min(i + save_batch_size, length)
|
||||
memmap[i: j] = corpus_embeddings[i: j]
|
||||
else:
|
||||
memmap[:] = corpus_embeddings
|
||||
|
||||
# create faiss index
|
||||
faiss_index = faiss.index_factory(dim, index_factory, faiss.METRIC_INNER_PRODUCT)
|
||||
|
||||
|
||||
if model.device == torch.device("cuda"):
|
||||
# co = faiss.GpuClonerOptions()
|
||||
co = faiss.GpuMultipleClonerOptions()
|
||||
co.useFloat16 = True
|
||||
# faiss_index = faiss.index_cpu_to_gpu(faiss.StandardGpuResources(), 0, faiss_index, co)
|
||||
faiss_index = faiss.index_cpu_to_all_gpus(faiss_index, co)
|
||||
|
||||
# NOTE: faiss only accepts float32
|
||||
logger.info("Adding embeddings...")
|
||||
corpus_embeddings = corpus_embeddings.astype(np.float32)
|
||||
faiss_index.train(corpus_embeddings)
|
||||
faiss_index.add(corpus_embeddings)
|
||||
|
||||
|
||||
|
||||
return faiss_index
|
||||
|
||||
|
||||
def search(model: Flag_mmret, queries: datasets, faiss_index: faiss.Index, k:int = 100, batch_size: int = 256, max_length: int=512):
|
||||
"""
|
||||
1. Encode queries into dense embeddings;
|
||||
2. Search through faiss index
|
||||
"""
|
||||
query_embeddings = model.encode_queries([queries["q_text"], queries["q_img"]],
|
||||
batch_size=batch_size,
|
||||
max_length=max_length,
|
||||
query_type='mm_it')
|
||||
|
||||
query_size = len(query_embeddings)
|
||||
|
||||
all_scores = []
|
||||
all_indices = []
|
||||
|
||||
for i in tqdm(range(0, query_size, batch_size), desc="Searching"):
|
||||
j = min(i + batch_size, query_size)
|
||||
query_embedding = query_embeddings[i: j]
|
||||
score, indice = faiss_index.search(query_embedding.astype(np.float32), k=k)
|
||||
all_scores.append(score)
|
||||
all_indices.append(indice)
|
||||
|
||||
all_scores = np.concatenate(all_scores, axis=0)
|
||||
all_indices = np.concatenate(all_indices, axis=0)
|
||||
return all_scores, all_indices
|
||||
|
||||
|
||||
def evaluate(preds, labels, cutoffs=[1,5,10,20,50,100]):
|
||||
"""
|
||||
Evaluate MRR and Recall at cutoffs.
|
||||
"""
|
||||
metrics = {}
|
||||
|
||||
# MRR
|
||||
mrrs = np.zeros(len(cutoffs))
|
||||
for pred, label in zip(preds, labels):
|
||||
jump = False
|
||||
for i, x in enumerate(pred, 1):
|
||||
if x in label:
|
||||
for k, cutoff in enumerate(cutoffs):
|
||||
if i <= cutoff:
|
||||
mrrs[k] += 1 / i
|
||||
jump = True
|
||||
if jump:
|
||||
break
|
||||
mrrs /= len(preds)
|
||||
for i, cutoff in enumerate(cutoffs):
|
||||
mrr = mrrs[i]
|
||||
metrics[f"MRR@{cutoff}"] = mrr
|
||||
|
||||
# Recall
|
||||
recalls = np.zeros(len(cutoffs))
|
||||
for pred, label in zip(preds, labels):
|
||||
if not isinstance(label, list):
|
||||
label = [label]
|
||||
for k, cutoff in enumerate(cutoffs):
|
||||
recall = np.intersect1d(label, pred[:cutoff])
|
||||
recalls[k] += len(recall) / len(label)
|
||||
recalls /= len(preds)
|
||||
for i, cutoff in enumerate(cutoffs):
|
||||
recall = recalls[i]
|
||||
metrics[f"Recall@{cutoff}"] = recall
|
||||
|
||||
return metrics
|
||||
|
||||
def main():
|
||||
parser = HfArgumentParser([Args])
|
||||
args: Args = parser.parse_args_into_dataclasses()[0]
|
||||
|
||||
model = Flag_mmret(model_name=args.model_name,
|
||||
normlized = True,
|
||||
image_dir=args.image_dir,
|
||||
use_fp16=False,
|
||||
)
|
||||
|
||||
eval_data = datasets.load_dataset('json', data_files="./eval/data/fashioniq_shirt_query_val.jsonl", split='train')
|
||||
image_corpus = datasets.load_dataset('json', data_files="./eval/data/fashioniq_shirt_corpus.jsonl", split='train')
|
||||
|
||||
faiss_index = index(
|
||||
model=model,
|
||||
corpus=image_corpus,
|
||||
batch_size=args.batch_size,
|
||||
max_length=args.max_passage_length,
|
||||
index_factory=args.index_factory,
|
||||
save_path=args.save_path,
|
||||
save_embedding=args.save_embedding,
|
||||
load_embedding=args.load_embedding
|
||||
)
|
||||
|
||||
scores, indices = search(
|
||||
model=model,
|
||||
queries=eval_data,
|
||||
faiss_index=faiss_index,
|
||||
k=args.k,
|
||||
batch_size=args.batch_size,
|
||||
max_length=args.max_query_length
|
||||
)
|
||||
|
||||
|
||||
|
||||
retrieval_results = []
|
||||
for indice in indices:
|
||||
# filter invalid indices
|
||||
indice = indice[indice != -1].tolist()
|
||||
retrieval_results.append(image_corpus[indice]["content"])
|
||||
|
||||
ground_truths = []
|
||||
for sample in eval_data:
|
||||
ground_truths.append(sample["positive_key"])
|
||||
|
||||
metrics_shirt = evaluate(retrieval_results, ground_truths)
|
||||
print("FashionIQ tasks (shirt):")
|
||||
print(metrics_shirt)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
eval_data = datasets.load_dataset('json', data_files="./eval/data/fashioniq_dress_query_val.jsonl", split='train')
|
||||
image_corpus = datasets.load_dataset('json', data_files="./eval/data/fashioniq_dress_corpus.jsonl", split='train')
|
||||
|
||||
faiss_index = index(
|
||||
model=model,
|
||||
corpus=image_corpus,
|
||||
batch_size=args.batch_size,
|
||||
max_length=args.max_passage_length,
|
||||
index_factory=args.index_factory,
|
||||
save_path=args.save_path,
|
||||
save_embedding=args.save_embedding,
|
||||
load_embedding=args.load_embedding
|
||||
)
|
||||
|
||||
scores, indices = search(
|
||||
model=model,
|
||||
queries=eval_data,
|
||||
faiss_index=faiss_index,
|
||||
k=args.k,
|
||||
batch_size=args.batch_size,
|
||||
max_length=args.max_query_length
|
||||
)
|
||||
|
||||
|
||||
|
||||
retrieval_results = []
|
||||
for indice in indices:
|
||||
# filter invalid indices
|
||||
indice = indice[indice != -1].tolist()
|
||||
retrieval_results.append(image_corpus[indice]["content"])
|
||||
|
||||
ground_truths = []
|
||||
for sample in eval_data:
|
||||
ground_truths.append(sample["positive_key"])
|
||||
|
||||
metrics_dress = evaluate(retrieval_results, ground_truths)
|
||||
print("FashionIQ tasks (dress):")
|
||||
print(metrics_dress)
|
||||
|
||||
|
||||
eval_data = datasets.load_dataset('json', data_files="./eval/data/fashioniq_toptee_query_val.jsonl", split='train')
|
||||
image_corpus = datasets.load_dataset('json', data_files="./eval/data/fashioniq_toptee_corpus.jsonl", split='train')
|
||||
|
||||
|
||||
faiss_index = index(
|
||||
model=model,
|
||||
corpus=image_corpus,
|
||||
batch_size=args.batch_size,
|
||||
max_length=args.max_passage_length,
|
||||
index_factory=args.index_factory,
|
||||
save_path=args.save_path,
|
||||
save_embedding=args.save_embedding,
|
||||
load_embedding=args.load_embedding
|
||||
)
|
||||
|
||||
scores, indices = search(
|
||||
model=model,
|
||||
queries=eval_data,
|
||||
faiss_index=faiss_index,
|
||||
k=args.k,
|
||||
batch_size=args.batch_size,
|
||||
max_length=args.max_query_length
|
||||
)
|
||||
|
||||
|
||||
|
||||
retrieval_results = []
|
||||
for indice in indices:
|
||||
# filter invalid indices
|
||||
indice = indice[indice != -1].tolist()
|
||||
retrieval_results.append(image_corpus[indice]["content"])
|
||||
|
||||
ground_truths = []
|
||||
for sample in eval_data:
|
||||
ground_truths.append(sample["positive_key"])
|
||||
|
||||
metrics_toptee = evaluate(retrieval_results, ground_truths)
|
||||
print("FashionIQ tasks (toptee):")
|
||||
print(metrics_toptee)
|
||||
|
||||
|
||||
|
||||
print(f"shirt: {metrics_shirt['Recall@10'] * 100:.2f} / {metrics_shirt['Recall@50'] * 100:.2f}")
|
||||
print(f"dress: {metrics_dress['Recall@10'] * 100:.2f} / {metrics_dress['Recall@50'] * 100:.2f}")
|
||||
print(f"toptee: {metrics_toptee['Recall@10'] * 100:.2f} / {metrics_toptee['Recall@50'] * 100:.2f}")
|
||||
print(f"overall: {(metrics_shirt['Recall@10'] + metrics_dress['Recall@10'] + metrics_toptee['Recall@10']) * 100 / 3:.2f} / {(metrics_shirt['Recall@50'] + metrics_dress['Recall@50'] + metrics_toptee['Recall@50']) * 100 / 3:.2f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,109 @@
|
||||
import math
|
||||
import os.path
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterator
|
||||
|
||||
import datasets
|
||||
from torch.utils.data import Dataset, IterableDataset
|
||||
from transformers import DataCollatorWithPadding
|
||||
from transformers import PreTrainedTokenizer, BatchEncoding
|
||||
from transformers import CLIPImageProcessor
|
||||
|
||||
|
||||
from PIL import Image
|
||||
import json
|
||||
import torch
|
||||
import torch.distributed
|
||||
|
||||
from io import BytesIO
|
||||
import warnings
|
||||
|
||||
class MMIT_Dataset(Dataset):
|
||||
def __init__(self, captions, image_ids, image_dir, image_processor) -> None:
|
||||
img_id_example = image_ids[0]
|
||||
img_id_example = str(img_id_example)
|
||||
if img_id_example[-4:] in [".jpg", ".png", "JPEG"]:
|
||||
self.image_path =[os.path.join(image_dir, str(id)) for id in image_ids]
|
||||
else:
|
||||
warnings.warn("Not found file extention in image_ids, will forcefully add '.jpg'.", UserWarning)
|
||||
self.image_path =[os.path.join(image_dir, str(id) + '.jpg') for id in image_ids]
|
||||
self.captions = captions
|
||||
self.image_processor = image_processor
|
||||
|
||||
def __getitem__(self, item):
|
||||
pil_data = Image.open(self.image_path[item])
|
||||
pil_data = pil_data.convert('RGB')
|
||||
image = self.image_processor(pil_data)
|
||||
|
||||
|
||||
|
||||
|
||||
caption = self.captions[item]
|
||||
|
||||
return caption, image
|
||||
|
||||
def __len__(self):
|
||||
return len(self.image_path)
|
||||
|
||||
|
||||
class MMIT_Collator:
|
||||
def __init__(self, tokenizer, caption_max_len):
|
||||
self.tokenizer = tokenizer
|
||||
self.caption_max_len = caption_max_len
|
||||
|
||||
|
||||
|
||||
def __call__(self, features):
|
||||
caption = [f[0] for f in features]
|
||||
images = [f[1] for f in features]
|
||||
|
||||
c_collated = self.tokenizer(
|
||||
caption,
|
||||
truncation=True,
|
||||
padding = True,
|
||||
max_length=self.caption_max_len,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
# i_collated = torch.stack(images)
|
||||
|
||||
# for clip model
|
||||
images = [f["pixel_values"][0] for f in images]
|
||||
images = [torch.tensor(arr) for arr in images]
|
||||
i_collated = torch.stack(images)
|
||||
##clip_end
|
||||
|
||||
return c_collated, i_collated
|
||||
|
||||
class Image_Dataset(Dataset):
|
||||
def __init__(self, image_ids, image_dir, image_processor) -> None:
|
||||
|
||||
self.image_path =[os.path.join(image_dir, str(id)) for id in image_ids]
|
||||
self.image_processor = image_processor
|
||||
|
||||
def __getitem__(self, item):
|
||||
pil_data = Image.open(self.image_path[item])
|
||||
image = self.image_processor(pil_data)
|
||||
|
||||
return image
|
||||
|
||||
def __len__(self):
|
||||
return len(self.image_path)
|
||||
|
||||
class Image_Collator:
|
||||
def __init__(self, tokenizer, caption_max_len):
|
||||
self.tokenizer = tokenizer
|
||||
self.caption_max_len = caption_max_len
|
||||
|
||||
|
||||
def __call__(self, features):
|
||||
# images = features
|
||||
# i_collated = torch.stack(images)
|
||||
|
||||
# for clip model
|
||||
images = [f["pixel_values"][0] for f in features]
|
||||
images = [torch.tensor(arr) for arr in images]
|
||||
i_collated = torch.stack(images)
|
||||
## clip-end
|
||||
return i_collated
|
||||
@@ -0,0 +1,206 @@
|
||||
from typing import cast, List, Union, Tuple
|
||||
import numpy as np
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoModel, AutoTokenizer, AutoModelForSequenceClassification, CLIPModel, CLIPImageProcessor, CLIPTokenizer
|
||||
import os
|
||||
from PIL import Image
|
||||
from torch.utils.data import DataLoader
|
||||
from torch import nn
|
||||
from flag_dataset import MMIT_Dataset, MMIT_Collator, Image_Dataset, Image_Collator
|
||||
class Flag_mmret(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str = None,
|
||||
normlized: bool = True,
|
||||
pooling_method: str = 'cls',
|
||||
use_fp16: bool=True,
|
||||
image_dir: str = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.model = AutoModel.from_pretrained(model_name)
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
self.image_processor = CLIPImageProcessor.from_pretrained(model_name)
|
||||
|
||||
|
||||
self.normalize_embeddings = normlized
|
||||
self.pooling_method = pooling_method
|
||||
|
||||
self.image_dir = image_dir
|
||||
|
||||
if use_fp16:
|
||||
self.use_fp16 = True
|
||||
self.model.half()
|
||||
else:
|
||||
self.use_fp16 = False
|
||||
|
||||
self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
|
||||
self.model = self.model.to(self.device)
|
||||
|
||||
self.num_gpus = torch.cuda.device_count()
|
||||
if self.num_gpus > 1:
|
||||
print(f"----------using {self.num_gpus}*GPUs----------")
|
||||
self.model = torch.nn.DataParallel(self.model)
|
||||
|
||||
|
||||
def encode_queries(self, queries: Union[List[str], str],
|
||||
batch_size: int=256,
|
||||
max_length: int=77,
|
||||
query_type: str = None,
|
||||
) -> np.ndarray:
|
||||
|
||||
|
||||
if query_type == 'text':
|
||||
input_texts = queries
|
||||
|
||||
return self.encode_text(input_texts, batch_size=batch_size, max_length=max_length)
|
||||
elif query_type == 'mm_it':
|
||||
q_text, q_img = queries
|
||||
|
||||
input_texts = q_text
|
||||
|
||||
|
||||
return self.encode_mm_it(input_texts, q_img, batch_size=batch_size)
|
||||
elif query_type == 'image':
|
||||
q_img = queries
|
||||
return self.encode_image(q_img, batch_size=batch_size)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def encode_corpus(self,
|
||||
corpus: dict,
|
||||
batch_size: int=256,
|
||||
max_length: int=77,
|
||||
corpus_type: str = None,
|
||||
) -> np.ndarray:
|
||||
if corpus_type == 'text':
|
||||
return self.encode_text(corpus["text"], batch_size=batch_size, max_length=max_length)
|
||||
elif corpus_type == 'mm_it':
|
||||
return self.encode_mm_it(corpus["text"], corpus["image"], batch_size=batch_size, max_length=max_length)
|
||||
elif corpus_type == 'image':
|
||||
return self.encode_image(corpus["image"], batch_size=batch_size, max_length=max_length)
|
||||
else:
|
||||
raise RuntimeError(f"You must choose a corpus type from: [mm_it, text, image]")
|
||||
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def encode_text(self, sentences: Union[List[str], str], batch_size: int=256, max_length: int=77) -> np.ndarray:
|
||||
if self.num_gpus > 0:
|
||||
batch_size = batch_size * self.num_gpus
|
||||
self.model.eval()
|
||||
|
||||
input_was_string = False
|
||||
if isinstance(sentences, str):
|
||||
sentences = [sentences]
|
||||
input_was_string = True
|
||||
|
||||
all_embeddings = []
|
||||
for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings", disable=len(sentences)<256):
|
||||
sentences_batch = sentences[start_index:start_index + batch_size]
|
||||
inputs = self.tokenizer(
|
||||
sentences_batch,
|
||||
padding=True,
|
||||
truncation=True,
|
||||
return_tensors='pt',
|
||||
max_length=max_length,
|
||||
).to(self.device)
|
||||
|
||||
embeddings = self.model.get_text_features(**inputs)
|
||||
embeddings = torch.nn.functional.normalize(embeddings, dim=-1)
|
||||
embeddings = cast(torch.Tensor, embeddings)
|
||||
all_embeddings.append(embeddings.cpu().numpy())
|
||||
|
||||
all_embeddings = np.concatenate(all_embeddings, axis=0)
|
||||
if input_was_string:
|
||||
return all_embeddings[0]
|
||||
return all_embeddings
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def encode_mm_it(self, captions: Union[List[str], str], image_ids: Union[List[str], str], batch_size: int=256, max_length: int=77) -> np.ndarray:
|
||||
if self.num_gpus > 0:
|
||||
batch_size = batch_size * self.num_gpus
|
||||
self.model.eval()
|
||||
|
||||
input_was_string = False
|
||||
if isinstance(captions, str):
|
||||
captions = [captions]
|
||||
image_ids = [image_ids]
|
||||
input_was_string = True
|
||||
|
||||
all_embeddings = []
|
||||
mm_it_dataset = MMIT_Dataset(captions=captions,
|
||||
image_ids=image_ids,
|
||||
image_dir=self.image_dir,
|
||||
image_processor=self.image_processor
|
||||
)
|
||||
mm_it_collator = MMIT_Collator(self.tokenizer, caption_max_len=75)
|
||||
|
||||
mm_it_dataloader = DataLoader(dataset=mm_it_dataset,
|
||||
collate_fn=mm_it_collator,
|
||||
num_workers=8,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
drop_last=False,)
|
||||
|
||||
for data in tqdm(mm_it_dataloader, desc="Inference Embeddings", disable=len(captions)<256):
|
||||
captions_inputs = data[0].to(self.device)
|
||||
|
||||
images = data[1].to(self.device)
|
||||
if self.use_fp16 and images.dtype != torch.float16:
|
||||
images = images.half()
|
||||
|
||||
text_embeddings = self.model.get_text_features(**captions_inputs)
|
||||
image_embeddings = self.model.get_image_features(images)
|
||||
|
||||
embeddings = text_embeddings + image_embeddings
|
||||
|
||||
embeddings = torch.nn.functional.normalize(embeddings, dim=-1)
|
||||
|
||||
embeddings = cast(torch.Tensor, embeddings)
|
||||
all_embeddings.append(embeddings.cpu().numpy())
|
||||
|
||||
all_embeddings = np.concatenate(all_embeddings, axis=0)
|
||||
if input_was_string:
|
||||
return all_embeddings[0]
|
||||
return all_embeddings
|
||||
|
||||
@torch.no_grad()
|
||||
def encode_image(self, image_ids: Union[List[str], str], batch_size: int=256, max_length: int=77) -> np.ndarray:
|
||||
if self.num_gpus > 0:
|
||||
batch_size = batch_size * self.num_gpus
|
||||
self.model.eval()
|
||||
|
||||
all_embeddings = []
|
||||
image_dataset = Image_Dataset(image_ids=image_ids,
|
||||
image_dir=self.image_dir,
|
||||
image_processor=self.image_processor
|
||||
)
|
||||
image_collator = Image_Collator(self.tokenizer, caption_max_len=312)
|
||||
|
||||
image_dataloader = DataLoader(dataset=image_dataset,
|
||||
collate_fn=image_collator,
|
||||
num_workers=8,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
drop_last=False,)
|
||||
|
||||
for data in tqdm(image_dataloader, desc="Inference Image Embeddings"):
|
||||
|
||||
images = data.to(self.device)
|
||||
if self.use_fp16 and images.dtype != torch.float16:
|
||||
images = images.half()
|
||||
|
||||
|
||||
|
||||
embeddings = self.model.get_image_features(images)
|
||||
embeddings = torch.nn.functional.normalize(embeddings, dim=-1)
|
||||
embeddings = cast(torch.Tensor, embeddings)
|
||||
all_embeddings.append(embeddings.cpu().numpy())
|
||||
|
||||
all_embeddings = np.concatenate(all_embeddings, axis=0)
|
||||
|
||||
return all_embeddings
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because it is too large
Load Diff
File diff suppressed because one or more lines are too long
@@ -0,0 +1,114 @@
|
||||
<h1 align="center">Vis-IR: Unifying Search With Visualized Information Retrieval</h1>
|
||||
|
||||
<p align="center">
|
||||
<a href="https://arxiv.org/abs/2502.11431">
|
||||
<img alt="Build" src="http://img.shields.io/badge/arXiv-2502.11431-B31B1B.svg">
|
||||
</a>
|
||||
<a href="https://github.com/VectorSpaceLab/Vis-IR">
|
||||
<img alt="Build" src="https://img.shields.io/badge/Github-Code-blue">
|
||||
</a>
|
||||
<a href="https://huggingface.co/datasets/marsh123/VIRA/">
|
||||
<img alt="Build" src="https://img.shields.io/badge/🤗 Datasets-VIRA-yellow">
|
||||
</a>
|
||||
<a href="https://huggingface.co/datasets/marsh123/MVRB">
|
||||
<img alt="Build" src="https://img.shields.io/badge/🤗 Datasets-MVRB-yellow">
|
||||
</a>
|
||||
<!-- <a href="">
|
||||
<img alt="Build" src="https://img.shields.io/badge/🤗 Model-UniSE CLIP-yellow">
|
||||
</a> -->
|
||||
<a href="https://huggingface.co/marsh123/UniSE">
|
||||
<img alt="Build" src="https://img.shields.io/badge/🤗 Model-UniSE MLLM-yellow">
|
||||
</a>
|
||||
|
||||
</p>
|
||||
<h4 align="center">
|
||||
<p>
|
||||
<a href=#news>News</a> |
|
||||
<a href=#release-plan>Release Plan</a> |
|
||||
<a href=#overview>Overview</a> |
|
||||
<a href="#license">License</a> |
|
||||
<a href="#citation">Citation</a>
|
||||
<p>
|
||||
</h4>
|
||||
|
||||
## News
|
||||
|
||||
```2025-04-06``` 🚀🚀 MVRB Dataset are released on Huggingface: [MVRB](https://huggingface.co/datasets/marsh123/MVRB)
|
||||
|
||||
```2025-04-02``` 🚀🚀 VIRA Dataset are released on Huggingface: [VIRA](https://huggingface.co/datasets/marsh123/VIRA/)
|
||||
|
||||
```2025-04-01``` 🚀🚀 UniSE models are released on Huggingface: [UniSE-MLMM](https://huggingface.co/marsh123/UniSE-MLLM/)
|
||||
|
||||
```2025-02-17``` 🎉🎉 Release our paper: [Any Information Is Just Worth One Single Screenshot: Unifying Search With Visualized Information Retrieval](https://arxiv.org/abs/2502.11431).
|
||||
|
||||
## Release Plan
|
||||
- [x] Paper
|
||||
- [x] UniSE models
|
||||
- [x] VIRA Dataset
|
||||
- [x] MVRB benchmark
|
||||
- [ ] Evaluation code
|
||||
- [ ] Fine-tuning code
|
||||
|
||||
## Overview
|
||||
|
||||
In this work, we formally define an emerging IR paradigm called Visualized Information Retrieval, or **VisIR**, where multimodal information, such as texts, images, tables and charts, is jointly represented by a unified visual format called **Screenshots**, for various retrieval applications. We further make three key contributions for VisIR. First, we create **VIRA** (Vis-IR Aggregation), a large-scale dataset comprising a vast collection of screenshots from diverse sources, carefully curated into captioned and questionanswer formats. Second, we develop **UniSE** (Universal Screenshot Embeddings), a family of retrieval models that enable screenshots to query or be queried across arbitrary data modalities. Finally, we construct **MVRB** (Massive Visualized IR Benchmark), a comprehensive benchmark covering a variety of task forms and application scenarios. Through extensive evaluations on MVRB, we highlight the deficiency from existing multimodal retrievers and the substantial improvements made by UniSE.
|
||||
|
||||
## Model Usage
|
||||
|
||||
> Our code works well on transformers==4.45.2, and we recommend using this version.
|
||||
|
||||
### 1. UniSE-MLLM Models
|
||||
|
||||
```python
|
||||
import torch
|
||||
from transformers import AutoModel
|
||||
|
||||
MODEL_NAME = "marsh123/UniSE-MLLM"
|
||||
model = AutoModel.from_pretrained(MODEL_NAME, trust_remote_code=True)
|
||||
# You must set trust_remote_code=True
|
||||
model.set_processor(MODEL_NAME)
|
||||
|
||||
with torch.no_grad():
|
||||
device = torch.device("cuda:0")
|
||||
model = model.to(device)
|
||||
model.eval()
|
||||
query_inputs = model.data_process(
|
||||
images=["./assets/query_1.png", "./assets/query_2.png"],
|
||||
text=["After a 17% drop, what is Nvidia's closing stock price?",
|
||||
"I would like to see a detailed and intuitive performance comparison between the two models."],
|
||||
q_or_c="query",
|
||||
task_instruction="Represent the given image with the given query."
|
||||
)
|
||||
candidate_inputs = model.data_process(
|
||||
images=["./assets/positive_1.jpeg", "./assets/neg_1.jpeg",
|
||||
"./assets/positive_2.jpeg", "./assets/neg_2.jpeg"],
|
||||
q_or_c="candidate"
|
||||
)
|
||||
query_embeddings = model(**query_inputs)
|
||||
candidate_embeddings = model(**candidate_inputs)
|
||||
scores = torch.matmul(query_embeddings, candidate_embeddings.T)
|
||||
print(scores)
|
||||
```
|
||||
|
||||
## Performance on MVRB
|
||||
|
||||
MVRB is a comprehensive benchmark designed for the retrieval task centered on screenshots. It includes four meta tasks: Screenshot Retrieval (SR), Composed Screenshot Retrieval (CSR), Screenshot QA (SQA), and Open-Vocabulary Classification (OVC). We evaluate three main types of retrievers on MVRB: OCR+Text Retrievers, General Multimodal Retrievers, and Screenshot Document Retrievers. Our proposed UniSE-MLLM achieves state-of-the-art (SOTA) performance on this benchmark.
|
||||
|
||||
|
||||
|
||||
|
||||
## License
|
||||
Vis-IR is licensed under the [MIT License](LICENSE).
|
||||
|
||||
|
||||
## Citation
|
||||
If you find this repository useful, please consider giving a star ⭐ and citation
|
||||
|
||||
```
|
||||
@article{liu2025any,
|
||||
title={Any Information Is Just Worth One Single Screenshot: Unifying Search With Visualized Information Retrieval},
|
||||
author={Liu, Ze and Liang, Zhengyang and Zhou, Junjie and Liu, Zheng and Lian, Defu},
|
||||
journal={arXiv preprint arXiv:2502.11431},
|
||||
year={2025}
|
||||
}
|
||||
```
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 2.0 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.4 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.2 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 649 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 234 KiB |
@@ -0,0 +1,11 @@
|
||||
# from .tasks import *
|
||||
from .tasks import *
|
||||
|
||||
ChineseTaskList = [
|
||||
'TNews', 'IFlyTek', 'MultilingualSentiment', 'JDReview', 'OnlineShopping', 'Waimai',
|
||||
'CLSClusteringS2S.v2', 'CLSClusteringP2P.v2', 'ThuNewsClusteringS2S.v2', 'ThuNewsClusteringP2P.v2',
|
||||
'Ocnli', 'Cmnli',
|
||||
'T2Reranking', 'MMarcoReranking', 'CMedQAv1-reranking', 'CMedQAv2-reranking',
|
||||
'T2Retrieval', 'MMarcoRetrieval', 'DuRetrieval', 'CovidRetrieval', 'CmedqaRetrieval', 'EcomRetrieval', 'MedicalRetrieval', 'VideoRetrieval',
|
||||
'ATEC', 'BQ', 'LCQMC', 'PAWSX', 'STSB', 'AFQMC', 'QBQTC'
|
||||
]
|
||||
@@ -0,0 +1,304 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from mteb.abstasks.AbsTaskClassification import AbsTaskClassification
|
||||
from mteb.abstasks.TaskMetadata import TaskMetadata
|
||||
|
||||
|
||||
class TNews(AbsTaskClassification):
|
||||
metadata = TaskMetadata(
|
||||
name="TNews",
|
||||
description="Short Text Classification for News",
|
||||
reference="https://www.cluebenchmarks.com/introduce.html",
|
||||
dataset={
|
||||
"path": "C-MTEB/TNews-classification",
|
||||
"revision": "317f262bf1e6126357bbe89e875451e4b0938fe4",
|
||||
},
|
||||
type="Classification",
|
||||
category="s2s",
|
||||
modalities=["text"],
|
||||
eval_splits=["validation"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="accuracy",
|
||||
date=None,
|
||||
domains=None,
|
||||
task_subtypes=None,
|
||||
license=None,
|
||||
annotations_creators=None,
|
||||
dialect=None,
|
||||
sample_creation=None,
|
||||
bibtex_citation="""@inproceedings {xu-etal-2020-clue,
|
||||
title = "{CLUE}: A {C}hinese Language Understanding Evaluation Benchmark",
|
||||
author = "Xu, Liang and
|
||||
Hu, Hai and
|
||||
Zhang, Xuanwei and
|
||||
Li, Lu and
|
||||
Cao, Chenjie and
|
||||
Li, Yudong and
|
||||
Xu, Yechen and
|
||||
Sun, Kai and
|
||||
Yu, Dian and
|
||||
Yu, Cong and
|
||||
Tian, Yin and
|
||||
Dong, Qianqian and
|
||||
Liu, Weitang and
|
||||
Shi, Bo and
|
||||
Cui, Yiming and
|
||||
Li, Junyi and
|
||||
Zeng, Jun and
|
||||
Wang, Rongzhao and
|
||||
Xie, Weijian and
|
||||
Li, Yanting and
|
||||
Patterson, Yina and
|
||||
Tian, Zuoyu and
|
||||
Zhang, Yiwen and
|
||||
Zhou, He and
|
||||
Liu, Shaoweihua and
|
||||
Zhao, Zhe and
|
||||
Zhao, Qipeng and
|
||||
Yue, Cong and
|
||||
Zhang, Xinrui and
|
||||
Yang, Zhengliang and
|
||||
Richardson, Kyle and
|
||||
Lan, Zhenzhong ",
|
||||
booktitle = "Proceedings of the 28th International Conference on Computational Linguistics",
|
||||
month = dec,
|
||||
year = "2020",
|
||||
address = "Barcelona, Spain (Online)",
|
||||
publisher = "International Committee on Computational Linguistics",
|
||||
url = "https://aclanthology.org/2020.coling-main.419",
|
||||
doi = "10.18653/v1/2020.coling-main.419",
|
||||
pages = "4762--4772",
|
||||
}""",
|
||||
descriptive_stats={"n_samples": None, "avg_character_length": None},
|
||||
)
|
||||
|
||||
@property
|
||||
def metadata_dict(self) -> dict[str, str]:
|
||||
metadata_dict = super().metadata_dict
|
||||
metadata_dict["samples_per_label"] = 32
|
||||
return metadata_dict
|
||||
|
||||
|
||||
class IFlyTek(AbsTaskClassification):
|
||||
metadata = TaskMetadata(
|
||||
name="IFlyTek",
|
||||
description="Long Text classification for the description of Apps",
|
||||
reference="https://www.cluebenchmarks.com/introduce.html",
|
||||
dataset={
|
||||
"path": "C-MTEB/IFlyTek-classification",
|
||||
"revision": "421605374b29664c5fc098418fe20ada9bd55f8a",
|
||||
},
|
||||
type="Classification",
|
||||
category="s2s",
|
||||
modalities=["text"],
|
||||
eval_splits=["validation"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="accuracy",
|
||||
date=None,
|
||||
domains=None,
|
||||
task_subtypes=None,
|
||||
license=None,
|
||||
annotations_creators=None,
|
||||
dialect=None,
|
||||
sample_creation=None,
|
||||
bibtex_citation="""@inproceedings {xu-etal-2020-clue,
|
||||
title = "{CLUE}: A {C}hinese Language Understanding Evaluation Benchmark",
|
||||
author = "Xu, Liang and
|
||||
Hu, Hai and
|
||||
Zhang, Xuanwei and
|
||||
Li, Lu and
|
||||
Cao, Chenjie and
|
||||
Li, Yudong and
|
||||
Xu, Yechen and
|
||||
Sun, Kai and
|
||||
Yu, Dian and
|
||||
Yu, Cong and
|
||||
Tian, Yin and
|
||||
Dong, Qianqian and
|
||||
Liu, Weitang and
|
||||
Shi, Bo and
|
||||
Cui, Yiming and
|
||||
Li, Junyi and
|
||||
Zeng, Jun and
|
||||
Wang, Rongzhao and
|
||||
Xie, Weijian and
|
||||
Li, Yanting and
|
||||
Patterson, Yina and
|
||||
Tian, Zuoyu and
|
||||
Zhang, Yiwen and
|
||||
Zhou, He and
|
||||
Liu, Shaoweihua and
|
||||
Zhao, Zhe and
|
||||
Zhao, Qipeng and
|
||||
Yue, Cong and
|
||||
Zhang, Xinrui and
|
||||
Yang, Zhengliang and
|
||||
Richardson, Kyle and
|
||||
Lan, Zhenzhong ",
|
||||
booktitle = "Proceedings of the 28th International Conference on Computational Linguistics",
|
||||
month = dec,
|
||||
year = "2020",
|
||||
address = "Barcelona, Spain (Online)",
|
||||
publisher = "International Committee on Computational Linguistics",
|
||||
url = "https://aclanthology.org/2020.coling-main.419",
|
||||
doi = "10.18653/v1/2020.coling-main.419",
|
||||
pages = "4762--4772",
|
||||
abstract = "The advent of natural language understanding (NLU) benchmarks for English, such as GLUE and SuperGLUE allows new NLU models to be evaluated across a diverse set of tasks. These comprehensive benchmarks have facilitated a broad range of research and applications in natural language processing (NLP). The problem, however, is that most such benchmarks are limited to English, which has made it difficult to replicate many of the successes in English NLU for other languages. To help remedy this issue, we introduce the first large-scale Chinese Language Understanding Evaluation (CLUE) benchmark. CLUE is an open-ended, community-driven project that brings together 9 tasks spanning several well-established single-sentence/sentence-pair classification tasks, as well as machine reading comprehension, all on original Chinese text. To establish results on these tasks, we report scores using an exhaustive set of current state-of-the-art pre-trained Chinese models (9 in total). We also introduce a number of supplementary datasets and additional tools to help facilitate further progress on Chinese NLU. Our benchmark is released at https://www.cluebenchmarks.com",
|
||||
}""",
|
||||
descriptive_stats={"n_samples": None, "avg_character_length": None},
|
||||
)
|
||||
|
||||
@property
|
||||
def metadata_dict(self) -> dict[str, str]:
|
||||
metadata_dict = super().metadata_dict
|
||||
metadata_dict["samples_per_label"] = 32
|
||||
metadata_dict["n_experiments"] = 5
|
||||
return metadata_dict
|
||||
|
||||
|
||||
class MultilingualSentiment(AbsTaskClassification):
|
||||
metadata = TaskMetadata(
|
||||
name="MultilingualSentiment",
|
||||
description="A collection of multilingual sentiments datasets grouped into 3 classes -- positive, neutral, negative",
|
||||
reference="https://github.com/tyqiangz/multilingual-sentiment-datasets",
|
||||
dataset={
|
||||
"path": "C-MTEB/MultilingualSentiment-classification",
|
||||
"revision": "46958b007a63fdbf239b7672c25d0bea67b5ea1a",
|
||||
},
|
||||
type="Classification",
|
||||
category="s2s",
|
||||
modalities=["text"],
|
||||
eval_splits=["validation", "test"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="accuracy",
|
||||
date=None,
|
||||
domains=None,
|
||||
task_subtypes=None,
|
||||
license=None,
|
||||
annotations_creators=None,
|
||||
dialect=None,
|
||||
sample_creation=None,
|
||||
bibtex_citation=None,
|
||||
descriptive_stats={"n_samples": None, "avg_character_length": None},
|
||||
)
|
||||
|
||||
@property
|
||||
def metadata_dict(self) -> dict[str, str]:
|
||||
metadata_dict = super().metadata_dict
|
||||
metadata_dict["samples_per_label"] = 32
|
||||
return metadata_dict
|
||||
|
||||
|
||||
class JDReview(AbsTaskClassification):
|
||||
metadata = TaskMetadata(
|
||||
name="JDReview",
|
||||
description="review for iphone",
|
||||
reference="https://aclanthology.org/2023.nodalida-1.20/",
|
||||
dataset={
|
||||
"path": "C-MTEB/JDReview-classification",
|
||||
"revision": "b7c64bd89eb87f8ded463478346f76731f07bf8b",
|
||||
},
|
||||
type="Classification",
|
||||
category="s2s",
|
||||
modalities=["text"],
|
||||
eval_splits=["test"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="accuracy",
|
||||
date=None,
|
||||
domains=None,
|
||||
task_subtypes=None,
|
||||
license=None,
|
||||
annotations_creators=None,
|
||||
dialect=None,
|
||||
sample_creation=None,
|
||||
bibtex_citation="""@article{xiao2023c,
|
||||
title={C-pack: Packaged resources to advance general chinese embedding},
|
||||
author={Xiao, Shitao and Liu, Zheng and Zhang, Peitian and Muennighof, Niklas},
|
||||
journal={arXiv preprint arXiv:2309.07597},
|
||||
year={2023}
|
||||
}""",
|
||||
descriptive_stats={"n_samples": None, "avg_character_length": None},
|
||||
)
|
||||
|
||||
@property
|
||||
def metadata_dict(self) -> dict[str, str]:
|
||||
metadata_dict = super().metadata_dict
|
||||
metadata_dict["samples_per_label"] = 32
|
||||
return metadata_dict
|
||||
|
||||
|
||||
class OnlineShopping(AbsTaskClassification):
|
||||
metadata = TaskMetadata(
|
||||
name="OnlineShopping",
|
||||
description="Sentiment Analysis of User Reviews on Online Shopping Websites",
|
||||
reference="https://aclanthology.org/2023.nodalida-1.20/",
|
||||
dataset={
|
||||
"path": "C-MTEB/OnlineShopping-classification",
|
||||
"revision": "e610f2ebd179a8fda30ae534c3878750a96db120",
|
||||
},
|
||||
type="Classification",
|
||||
category="s2s",
|
||||
modalities=["text"],
|
||||
eval_splits=["test"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="accuracy",
|
||||
date=None,
|
||||
domains=None,
|
||||
task_subtypes=None,
|
||||
license=None,
|
||||
annotations_creators=None,
|
||||
dialect=None,
|
||||
sample_creation=None,
|
||||
bibtex_citation="""@article{xiao2023c,
|
||||
title={C-pack: Packaged resources to advance general chinese embedding},
|
||||
author={Xiao, Shitao and Liu, Zheng and Zhang, Peitian and Muennighof, Niklas},
|
||||
journal={arXiv preprint arXiv:2309.07597},
|
||||
year={2023}
|
||||
}""",
|
||||
descriptive_stats={"n_samples": None, "avg_character_length": None},
|
||||
)
|
||||
|
||||
@property
|
||||
def metadata_dict(self) -> dict[str, str]:
|
||||
metadata_dict = super().metadata_dict
|
||||
metadata_dict["samples_per_label"] = 32
|
||||
return metadata_dict
|
||||
|
||||
|
||||
class Waimai(AbsTaskClassification):
|
||||
metadata = TaskMetadata(
|
||||
name="Waimai",
|
||||
description="Sentiment Analysis of user reviews on takeaway platforms",
|
||||
reference="https://aclanthology.org/2023.nodalida-1.20/",
|
||||
dataset={
|
||||
"path": "C-MTEB/waimai-classification",
|
||||
"revision": "339287def212450dcaa9df8c22bf93e9980c7023",
|
||||
},
|
||||
type="Classification",
|
||||
category="s2s",
|
||||
modalities=["text"],
|
||||
eval_splits=["test"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="accuracy",
|
||||
date=None,
|
||||
domains=None,
|
||||
task_subtypes=None,
|
||||
license=None,
|
||||
annotations_creators=None,
|
||||
dialect=None,
|
||||
sample_creation=None,
|
||||
bibtex_citation="""@article{xiao2023c,
|
||||
title={C-pack: Packaged resources to advance general chinese embedding},
|
||||
author={Xiao, Shitao and Liu, Zheng and Zhang, Peitian and Muennighof, Niklas},
|
||||
journal={arXiv preprint arXiv:2309.07597},
|
||||
year={2023}
|
||||
}""",
|
||||
descriptive_stats={"n_samples": None, "avg_character_length": None},
|
||||
)
|
||||
|
||||
@property
|
||||
def metadata_dict(self) -> dict[str, str]:
|
||||
metadata_dict = super().metadata_dict
|
||||
metadata_dict["samples_per_label"] = 32
|
||||
|
||||
return metadata_dict
|
||||
@@ -0,0 +1,410 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
|
||||
from datasets import Dataset, DatasetDict
|
||||
|
||||
from mteb.abstasks.AbsTaskClustering import AbsTaskClustering
|
||||
from mteb.abstasks.AbsTaskClusteringFast import (
|
||||
AbsTaskClusteringFast,
|
||||
check_label_distribution,
|
||||
)
|
||||
from mteb.abstasks.TaskMetadata import TaskMetadata
|
||||
|
||||
NUM_SAMPLES = 2048
|
||||
|
||||
|
||||
class CLSClusteringFastS2S(AbsTaskClusteringFast):
|
||||
max_document_to_embed = NUM_SAMPLES
|
||||
max_fraction_of_documents_to_embed = None
|
||||
|
||||
metadata = TaskMetadata(
|
||||
name="CLSClusteringS2S.v2",
|
||||
description="Clustering of titles from CLS dataset. Clustering of 13 sets on the main category.",
|
||||
reference="https://arxiv.org/abs/2209.05034",
|
||||
dataset={
|
||||
"path": "C-MTEB/CLSClusteringS2S",
|
||||
"revision": "e458b3f5414b62b7f9f83499ac1f5497ae2e869f",
|
||||
},
|
||||
type="Clustering",
|
||||
category="s2s",
|
||||
modalities=["text"],
|
||||
eval_splits=["test"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="v_measure",
|
||||
date=("2022-01-01", "2022-09-12"),
|
||||
domains=["Academic", "Written"],
|
||||
task_subtypes=["Thematic clustering", "Topic classification"],
|
||||
license="apache-2.0",
|
||||
annotations_creators="derived",
|
||||
dialect=[],
|
||||
sample_creation="found",
|
||||
bibtex_citation="""@misc{li2022csl,
|
||||
title={CSL: A Large-scale Chinese Scientific Literature Dataset},
|
||||
author={Yudong Li and Yuqing Zhang and Zhe Zhao and Linlin Shen and Weijie Liu and Weiquan Mao and Hui Zhang},
|
||||
year={2022},
|
||||
eprint={2209.05034},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CL}
|
||||
}""",
|
||||
descriptive_stats={
|
||||
"n_samples": {"test": NUM_SAMPLES},
|
||||
"avg_character_length": {},
|
||||
},
|
||||
)
|
||||
|
||||
def dataset_transform(self):
|
||||
ds = {}
|
||||
for split in self.metadata.eval_splits:
|
||||
labels = list(itertools.chain.from_iterable(self.dataset[split]["labels"]))
|
||||
sentences = list(
|
||||
itertools.chain.from_iterable(self.dataset[split]["sentences"])
|
||||
)
|
||||
|
||||
check_label_distribution(self.dataset[split])
|
||||
|
||||
ds[split] = Dataset.from_dict({"labels": labels, "sentences": sentences})
|
||||
self.dataset = DatasetDict(ds)
|
||||
self.dataset = self.stratified_subsampling(
|
||||
self.dataset,
|
||||
self.seed,
|
||||
self.metadata.eval_splits,
|
||||
label="labels",
|
||||
n_samples=NUM_SAMPLES,
|
||||
)
|
||||
|
||||
|
||||
class CLSClusteringFastP2P(AbsTaskClusteringFast):
|
||||
max_document_to_embed = NUM_SAMPLES
|
||||
max_fraction_of_documents_to_embed = None
|
||||
|
||||
metadata = TaskMetadata(
|
||||
name="CLSClusteringP2P.v2",
|
||||
description="Clustering of titles + abstract from CLS dataset. Clustering of 13 sets on the main category.",
|
||||
reference="https://arxiv.org/abs/2209.05034",
|
||||
dataset={
|
||||
"path": "C-MTEB/CLSClusteringP2P",
|
||||
"revision": "4b6227591c6c1a73bc76b1055f3b7f3588e72476",
|
||||
},
|
||||
type="Clustering",
|
||||
category="p2p",
|
||||
modalities=["text"],
|
||||
eval_splits=["test"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="v_measure",
|
||||
date=("2022-01-01", "2022-09-12"),
|
||||
domains=["Academic", "Written"],
|
||||
task_subtypes=["Thematic clustering", "Topic classification"],
|
||||
license="apache-2.0",
|
||||
annotations_creators="derived",
|
||||
dialect=[],
|
||||
sample_creation="found",
|
||||
bibtex_citation="""@misc{li2022csl,
|
||||
title={CSL: A Large-scale Chinese Scientific Literature Dataset},
|
||||
author={Yudong Li and Yuqing Zhang and Zhe Zhao and Linlin Shen and Weijie Liu and Weiquan Mao and Hui Zhang},
|
||||
year={2022},
|
||||
eprint={2209.05034},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CL}
|
||||
}""",
|
||||
descriptive_stats={
|
||||
"n_samples": {"test": NUM_SAMPLES},
|
||||
"avg_character_length": {},
|
||||
},
|
||||
)
|
||||
|
||||
def dataset_transform(self):
|
||||
ds = {}
|
||||
for split in self.metadata.eval_splits:
|
||||
labels = list(itertools.chain.from_iterable(self.dataset[split]["labels"]))
|
||||
sentences = list(
|
||||
itertools.chain.from_iterable(self.dataset[split]["sentences"])
|
||||
)
|
||||
|
||||
check_label_distribution(self.dataset[split])
|
||||
|
||||
ds[split] = Dataset.from_dict({"labels": labels, "sentences": sentences})
|
||||
self.dataset = DatasetDict(ds)
|
||||
self.dataset = self.stratified_subsampling(
|
||||
self.dataset,
|
||||
self.seed,
|
||||
self.metadata.eval_splits,
|
||||
label="labels",
|
||||
n_samples=NUM_SAMPLES,
|
||||
)
|
||||
|
||||
|
||||
class CLSClusteringS2S(AbsTaskClustering):
|
||||
superseded_by = "CLSClusteringS2S.v2"
|
||||
metadata = TaskMetadata(
|
||||
name="CLSClusteringS2S",
|
||||
description="Clustering of titles from CLS dataset. Clustering of 13 sets on the main category.",
|
||||
reference="https://arxiv.org/abs/2209.05034",
|
||||
dataset={
|
||||
"path": "C-MTEB/CLSClusteringS2S",
|
||||
"revision": "e458b3f5414b62b7f9f83499ac1f5497ae2e869f",
|
||||
},
|
||||
type="Clustering",
|
||||
category="s2s",
|
||||
modalities=["text"],
|
||||
eval_splits=["test"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="v_measure",
|
||||
date=None,
|
||||
form=None,
|
||||
domains=None,
|
||||
task_subtypes=None,
|
||||
license=None,
|
||||
annotations_creators=None,
|
||||
dialect=None,
|
||||
sample_creation=None,
|
||||
bibtex_citation="""
|
||||
@article{li2022csl,
|
||||
title={CSL: A large-scale Chinese scientific literature dataset},
|
||||
author={Li, Yudong and Zhang, Yuqing and Zhao, Zhe and Shen, Linlin and Liu, Weijie and Mao, Weiquan and Zhang, Hui},
|
||||
journal={arXiv preprint arXiv:2209.05034},
|
||||
year={2022}
|
||||
}
|
||||
""",
|
||||
descriptive_stats={"n_samples": {"test": 100000}, "avg_character_length": None},
|
||||
)
|
||||
|
||||
|
||||
class CLSClusteringP2P(AbsTaskClustering):
|
||||
superseded_by = "CLSClusteringP2P.v2"
|
||||
metadata = TaskMetadata(
|
||||
name="CLSClusteringP2P",
|
||||
description="Clustering of titles + abstract from CLS dataset. Clustering of 13 sets on the main category.",
|
||||
reference="https://arxiv.org/abs/2209.05034",
|
||||
dataset={
|
||||
"path": "C-MTEB/CLSClusteringP2P",
|
||||
"revision": "4b6227591c6c1a73bc76b1055f3b7f3588e72476",
|
||||
},
|
||||
type="Clustering",
|
||||
category="p2p",
|
||||
modalities=["text"],
|
||||
eval_splits=["test"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="v_measure",
|
||||
date=None,
|
||||
form=None,
|
||||
domains=None,
|
||||
task_subtypes=None,
|
||||
license=None,
|
||||
annotations_creators=None,
|
||||
dialect=None,
|
||||
sample_creation=None,
|
||||
bibtex_citation="""@article{li2022csl,
|
||||
title={CSL: A large-scale Chinese scientific literature dataset},
|
||||
author={Li, Yudong and Zhang, Yuqing and Zhao, Zhe and Shen, Linlin and Liu, Weijie and Mao, Weiquan and Zhang, Hui},
|
||||
journal={arXiv preprint arXiv:2209.05034},
|
||||
year={2022}
|
||||
}""",
|
||||
descriptive_stats={"n_samples": {"test": 100000}, "avg_character_length": None},
|
||||
)
|
||||
|
||||
|
||||
class ThuNewsClusteringFastS2S(AbsTaskClusteringFast):
|
||||
max_document_to_embed = NUM_SAMPLES
|
||||
max_fraction_of_documents_to_embed = None
|
||||
|
||||
metadata = TaskMetadata(
|
||||
name="ThuNewsClusteringS2S.v2",
|
||||
dataset={
|
||||
"path": "C-MTEB/ThuNewsClusteringS2S",
|
||||
"revision": "8a8b2caeda43f39e13c4bc5bea0f8a667896e10d",
|
||||
},
|
||||
description="Clustering of titles from the THUCNews dataset",
|
||||
reference="http://thuctc.thunlp.org/",
|
||||
type="Clustering",
|
||||
category="s2s",
|
||||
modalities=["text"],
|
||||
eval_splits=["test"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="v_measure",
|
||||
date=("2006-01-01", "2007-01-01"),
|
||||
domains=["News", "Written"],
|
||||
task_subtypes=["Thematic clustering", "Topic classification"],
|
||||
license="not specified",
|
||||
annotations_creators="derived",
|
||||
dialect=[],
|
||||
sample_creation="found",
|
||||
bibtex_citation="""@software{THUCTC,
|
||||
author = {Sun, M. and Li, J. and Guo, Z. and Yu, Z. and Zheng, Y. and Si, X. and Liu, Z.},
|
||||
title = {THUCTC: An Efficient Chinese Text Classifier},
|
||||
year = {2016},
|
||||
note = {THU Chinese Text Classification Toolkit},
|
||||
publisher = {THU Natural Language Processing Lab},
|
||||
url = {https://github.com/thunlp/THUCTC}
|
||||
}""",
|
||||
descriptive_stats={
|
||||
"n_samples": {"test": NUM_SAMPLES},
|
||||
"avg_character_length": {},
|
||||
},
|
||||
)
|
||||
|
||||
def dataset_transform(self):
|
||||
ds = {}
|
||||
for split in self.metadata.eval_splits:
|
||||
labels = list(itertools.chain.from_iterable(self.dataset[split]["labels"]))
|
||||
sentences = list(
|
||||
itertools.chain.from_iterable(self.dataset[split]["sentences"])
|
||||
)
|
||||
|
||||
check_label_distribution(self.dataset[split])
|
||||
|
||||
ds[split] = Dataset.from_dict({"labels": labels, "sentences": sentences})
|
||||
self.dataset = DatasetDict(ds)
|
||||
self.dataset = self.stratified_subsampling(
|
||||
self.dataset,
|
||||
self.seed,
|
||||
self.metadata.eval_splits,
|
||||
label="labels",
|
||||
n_samples=NUM_SAMPLES,
|
||||
)
|
||||
|
||||
|
||||
class ThuNewsClusteringFastP2P(AbsTaskClusteringFast):
|
||||
max_document_to_embed = NUM_SAMPLES
|
||||
max_fraction_of_documents_to_embed = None
|
||||
|
||||
metadata = TaskMetadata(
|
||||
name="ThuNewsClusteringP2P.v2",
|
||||
dataset={
|
||||
"path": "C-MTEB/ThuNewsClusteringP2P",
|
||||
"revision": "5798586b105c0434e4f0fe5e767abe619442cf93",
|
||||
},
|
||||
description="Clustering of titles + abstracts from the THUCNews dataset",
|
||||
reference="http://thuctc.thunlp.org/",
|
||||
type="Clustering",
|
||||
category="p2p",
|
||||
modalities=["text"],
|
||||
eval_splits=["test"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="v_measure",
|
||||
date=("2006-01-01", "2007-01-01"),
|
||||
domains=["News", "Written"],
|
||||
task_subtypes=["Thematic clustering", "Topic classification"],
|
||||
license="not specified",
|
||||
annotations_creators="derived",
|
||||
dialect=[],
|
||||
sample_creation="found",
|
||||
bibtex_citation="""@software{THUCTC,
|
||||
author = {Sun, M. and Li, J. and Guo, Z. and Yu, Z. and Zheng, Y. and Si, X. and Liu, Z.},
|
||||
title = {THUCTC: An Efficient Chinese Text Classifier},
|
||||
year = {2016},
|
||||
note = {THU Chinese Text Classification Toolkit},
|
||||
publisher = {THU Natural Language Processing Lab},
|
||||
url = {https://github.com/thunlp/THUCTC}
|
||||
}""",
|
||||
descriptive_stats={
|
||||
"n_samples": {"test": NUM_SAMPLES},
|
||||
"avg_character_length": {},
|
||||
},
|
||||
)
|
||||
|
||||
def dataset_transform(self):
|
||||
ds = {}
|
||||
for split in self.metadata.eval_splits:
|
||||
labels = list(itertools.chain.from_iterable(self.dataset[split]["labels"]))
|
||||
sentences = list(
|
||||
itertools.chain.from_iterable(self.dataset[split]["sentences"])
|
||||
)
|
||||
|
||||
check_label_distribution(self.dataset[split])
|
||||
|
||||
ds[split] = Dataset.from_dict({"labels": labels, "sentences": sentences})
|
||||
self.dataset = DatasetDict(ds)
|
||||
self.dataset = self.stratified_subsampling(
|
||||
self.dataset,
|
||||
self.seed,
|
||||
self.metadata.eval_splits,
|
||||
label="labels",
|
||||
n_samples=NUM_SAMPLES,
|
||||
)
|
||||
|
||||
|
||||
class ThuNewsClusteringS2S(AbsTaskClustering):
|
||||
superseded_by = "ThuNewsClusteringS2S.v2"
|
||||
metadata = TaskMetadata(
|
||||
name="ThuNewsClusteringS2S",
|
||||
dataset={
|
||||
"path": "C-MTEB/ThuNewsClusteringS2S",
|
||||
"revision": "8a8b2caeda43f39e13c4bc5bea0f8a667896e10d",
|
||||
},
|
||||
description="Clustering of titles from the THUCNews dataset",
|
||||
reference="http://thuctc.thunlp.org/",
|
||||
type="Clustering",
|
||||
category="s2s",
|
||||
modalities=["text"],
|
||||
eval_splits=["test"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="v_measure",
|
||||
date=None,
|
||||
form=None,
|
||||
domains=None,
|
||||
task_subtypes=None,
|
||||
license=None,
|
||||
annotations_creators=None,
|
||||
dialect=None,
|
||||
sample_creation=None,
|
||||
bibtex_citation="""
|
||||
@inproceedings{eisner2007proceedings,
|
||||
title={Proceedings of the 2007 joint conference on empirical methods in natural language processing and computational natural language learning (EMNLP-CoNLL)},
|
||||
author={Eisner, Jason},
|
||||
booktitle={Proceedings of the 2007 Joint Conference on Empirical Methods in Natural Language Processing and Computational Natural Language Learning (EMNLP-CoNLL)},
|
||||
year={2007}
|
||||
}
|
||||
@inproceedings{li2006comparison,
|
||||
title={A comparison and semi-quantitative analysis of words and character-bigrams as features in chinese text categorization},
|
||||
author={Li, Jingyang and Sun, Maosong and Zhang, Xian},
|
||||
booktitle={proceedings of the 21st international conference on computational linguistics and 44th annual meeting of the association for computational linguistics},
|
||||
pages={545--552},
|
||||
year={2006}
|
||||
}
|
||||
""",
|
||||
descriptive_stats={"n_samples": {"test": 100000}, "avg_character_length": None},
|
||||
)
|
||||
|
||||
|
||||
class ThuNewsClusteringP2P(AbsTaskClustering):
|
||||
superseded_by = "ThuNewsClusteringP2P.v2"
|
||||
metadata = TaskMetadata(
|
||||
name="ThuNewsClusteringP2P",
|
||||
dataset={
|
||||
"path": "C-MTEB/ThuNewsClusteringP2P",
|
||||
"revision": "5798586b105c0434e4f0fe5e767abe619442cf93",
|
||||
},
|
||||
description="Clustering of titles + abstracts from the THUCNews dataset",
|
||||
reference="http://thuctc.thunlp.org/",
|
||||
type="Clustering",
|
||||
category="p2p",
|
||||
modalities=["text"],
|
||||
eval_splits=["test"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="v_measure",
|
||||
date=None,
|
||||
form=None,
|
||||
domains=None,
|
||||
task_subtypes=None,
|
||||
license=None,
|
||||
annotations_creators=None,
|
||||
dialect=None,
|
||||
sample_creation=None,
|
||||
bibtex_citation="""
|
||||
@inproceedings{eisner2007proceedings,
|
||||
title={Proceedings of the 2007 joint conference on empirical methods in natural language processing and computational natural language learning (EMNLP-CoNLL)},
|
||||
author={Eisner, Jason},
|
||||
booktitle={Proceedings of the 2007 Joint Conference on Empirical Methods in Natural Language Processing and Computational Natural Language Learning (EMNLP-CoNLL)},
|
||||
year={2007}
|
||||
}
|
||||
@inproceedings{li2006comparison,
|
||||
title={A comparison and semi-quantitative analysis of words and character-bigrams as features in chinese text categorization},
|
||||
author={Li, Jingyang and Sun, Maosong and Zhang, Xian},
|
||||
booktitle={proceedings of the 21st international conference on computational linguistics and 44th annual meeting of the association for computational linguistics},
|
||||
pages={545--552},
|
||||
year={2006}
|
||||
}
|
||||
""",
|
||||
descriptive_stats={"n_samples": {"test": 100000}, "avg_character_length": None},
|
||||
)
|
||||
@@ -0,0 +1,117 @@
|
||||
import datasets
|
||||
from mteb.abstasks import MultilingualTask, AbsTaskRetrieval
|
||||
from mteb.abstasks.AbsTaskRetrieval import *
|
||||
# from ...abstasks import MultilingualTask, AbsTaskRetrieval
|
||||
# from ...abstasks.AbsTaskRetrieval import *
|
||||
|
||||
|
||||
_LANGUAGES = ['ar', 'de', 'en', 'es', 'fr', 'hi', 'it', 'ja', 'ko', 'pt', 'ru', 'th', 'zh']
|
||||
|
||||
|
||||
def load_mldr_data(path: str, langs: list, eval_splits: list, cache_dir: str=None):
|
||||
corpus = {lang: {split: None for split in eval_splits} for lang in langs}
|
||||
queries = {lang: {split: None for split in eval_splits} for lang in langs}
|
||||
relevant_docs = {lang: {split: None for split in eval_splits} for lang in langs}
|
||||
|
||||
for lang in langs:
|
||||
lang_corpus = datasets.load_dataset(path, f'corpus-{lang}', cache_dir=cache_dir)['corpus']
|
||||
lang_corpus = {e['docid']: {'text': e['text']} for e in lang_corpus}
|
||||
lang_data = datasets.load_dataset(path, lang, cache_dir=cache_dir)
|
||||
for split in eval_splits:
|
||||
corpus[lang][split] = lang_corpus
|
||||
queries[lang][split] = {e['query_id']: e['query'] for e in lang_data[split]}
|
||||
relevant_docs[lang][split] = {e['query_id']: {e['positive_passages'][0]['docid']: 1} for e in lang_data[split]}
|
||||
|
||||
corpus = datasets.DatasetDict(corpus)
|
||||
queries = datasets.DatasetDict(queries)
|
||||
relevant_docs = datasets.DatasetDict(relevant_docs)
|
||||
return corpus, queries, relevant_docs
|
||||
|
||||
|
||||
class MultiLongDocRetrieval(MultilingualTask, AbsTaskRetrieval):
|
||||
@property
|
||||
def description(self):
|
||||
return {
|
||||
'name': 'MultiLongDocRetrieval',
|
||||
'hf_hub_name': 'Shitao/MLDR',
|
||||
'reference': 'https://arxiv.org/abs/2402.03216',
|
||||
'description': 'MultiLongDocRetrieval: A Multilingual Long-Document Retrieval Dataset',
|
||||
'type': 'Retrieval',
|
||||
'category': 's2p',
|
||||
'eval_splits': ['dev', 'test'],
|
||||
'eval_langs': _LANGUAGES,
|
||||
'main_score': 'ndcg_at_10',
|
||||
}
|
||||
|
||||
def load_data(self, **kwargs):
|
||||
if self.data_loaded:
|
||||
return
|
||||
|
||||
self.corpus, self.queries, self.relevant_docs = load_mldr_data(
|
||||
path=self.description['hf_hub_name'],
|
||||
langs=self.langs,
|
||||
eval_splits=self.description['eval_splits'],
|
||||
cache_dir=kwargs.get('cache_dir', None)
|
||||
)
|
||||
self.data_loaded = True
|
||||
|
||||
def evaluate(
|
||||
self,
|
||||
model,
|
||||
split="test",
|
||||
batch_size=128,
|
||||
corpus_chunk_size=None,
|
||||
score_function="cos_sim",
|
||||
**kwargs
|
||||
):
|
||||
try:
|
||||
from beir.retrieval.evaluation import EvaluateRetrieval
|
||||
except ImportError:
|
||||
raise Exception("Retrieval tasks require beir package. Please install it with `pip install mteb[beir]`")
|
||||
|
||||
if not self.data_loaded:
|
||||
self.load_data()
|
||||
|
||||
model = model if self.is_dres_compatible(model) else DRESModel(model)
|
||||
if os.getenv("RANK", None) is None:
|
||||
# Non-distributed
|
||||
from beir.retrieval.search.dense import DenseRetrievalExactSearch as DRES
|
||||
model = DRES(
|
||||
model,
|
||||
batch_size=batch_size,
|
||||
corpus_chunk_size=corpus_chunk_size if corpus_chunk_size is not None else 50000,
|
||||
**kwargs,
|
||||
)
|
||||
else:
|
||||
# Distributed (multi-GPU)
|
||||
from beir.retrieval.search.dense import (
|
||||
DenseRetrievalParallelExactSearch as DRPES,
|
||||
)
|
||||
model = DRPES(
|
||||
model,
|
||||
batch_size=batch_size,
|
||||
corpus_chunk_size=corpus_chunk_size,
|
||||
**kwargs,
|
||||
)
|
||||
retriever = EvaluateRetrieval(model, score_function=score_function) # or "cos_sim" or "dot"
|
||||
|
||||
scores = {}
|
||||
for lang in self.langs:
|
||||
print(f"==============================\nStart evaluating {lang} ...")
|
||||
corpus, queries, relevant_docs = self.corpus[lang][split], self.queries[lang][split], self.relevant_docs[lang][split]
|
||||
|
||||
start_time = time()
|
||||
results = retriever.retrieve(corpus, queries)
|
||||
end_time = time()
|
||||
logger.info("Time taken to retrieve: {:.2f} seconds".format(end_time - start_time))
|
||||
ndcg, _map, recall, precision = retriever.evaluate(relevant_docs, results, retriever.k_values, ignore_identical_ids=kwargs.get("ignore_identical_ids", True))
|
||||
mrr = retriever.evaluate_custom(relevant_docs, results, retriever.k_values, "mrr")
|
||||
scores[lang] = {
|
||||
**{f"ndcg_at_{k.split('@')[1]}": v for (k, v) in ndcg.items()},
|
||||
**{f"map_at_{k.split('@')[1]}": v for (k, v) in _map.items()},
|
||||
**{f"recall_at_{k.split('@')[1]}": v for (k, v) in recall.items()},
|
||||
**{f"precision_at_{k.split('@')[1]}": v for (k, v) in precision.items()},
|
||||
**{f"mrr_at_{k.split('@')[1]}": v for (k, v) in mrr.items()},
|
||||
}
|
||||
|
||||
return scores
|
||||
@@ -0,0 +1,115 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from mteb.abstasks.AbsTaskPairClassification import AbsTaskPairClassification
|
||||
from mteb.abstasks.TaskMetadata import TaskMetadata
|
||||
|
||||
|
||||
class Ocnli(AbsTaskPairClassification):
|
||||
metadata = TaskMetadata(
|
||||
name="Ocnli",
|
||||
description="Original Chinese Natural Language Inference dataset",
|
||||
reference="https://arxiv.org/abs/2010.05444",
|
||||
dataset={
|
||||
"path": "C-MTEB/OCNLI",
|
||||
"revision": "66e76a618a34d6d565d5538088562851e6daa7ec",
|
||||
},
|
||||
type="PairClassification",
|
||||
category="s2s",
|
||||
modalities=["text"],
|
||||
eval_splits=["validation"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="max_accuracy",
|
||||
date=None,
|
||||
domains=None,
|
||||
task_subtypes=None,
|
||||
license=None,
|
||||
annotations_creators=None,
|
||||
dialect=None,
|
||||
sample_creation=None,
|
||||
bibtex_citation="""@misc{hu2020ocnli,
|
||||
title={OCNLI: Original Chinese Natural Language Inference},
|
||||
author={Hai Hu and Kyle Richardson and Liang Xu and Lu Li and Sandra Kuebler and Lawrence S. Moss},
|
||||
year={2020},
|
||||
eprint={2010.05444},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CL}
|
||||
}""",
|
||||
descriptive_stats={"n_samples": None, "avg_character_length": None},
|
||||
)
|
||||
|
||||
def dataset_transform(self):
|
||||
self.dataset = self.dataset.rename_column("sent1", "sentence1")
|
||||
self.dataset = self.dataset.rename_column("sent2", "sentence2")
|
||||
|
||||
|
||||
class Cmnli(AbsTaskPairClassification):
|
||||
metadata = TaskMetadata(
|
||||
name="Cmnli",
|
||||
description="Chinese Multi-Genre NLI",
|
||||
reference="https://huggingface.co/datasets/clue/viewer/cmnli",
|
||||
dataset={
|
||||
"path": "C-MTEB/CMNLI",
|
||||
"revision": "41bc36f332156f7adc9e38f53777c959b2ae9766",
|
||||
},
|
||||
type="PairClassification",
|
||||
category="s2s",
|
||||
modalities=["text"],
|
||||
eval_splits=["validation"],
|
||||
eval_langs=["cmn-Hans"],
|
||||
main_score="max_accuracy",
|
||||
date=None,
|
||||
domains=None,
|
||||
task_subtypes=None,
|
||||
license=None,
|
||||
annotations_creators=None,
|
||||
dialect=None,
|
||||
sample_creation=None,
|
||||
bibtex_citation="""@inproceedings{xu-etal-2020-clue,
|
||||
title = "{CLUE}: A {C}hinese Language Understanding Evaluation Benchmark",
|
||||
author = "Xu, Liang and
|
||||
Hu, Hai and
|
||||
Zhang, Xuanwei and
|
||||
Li, Lu and
|
||||
Cao, Chenjie and
|
||||
Li, Yudong and
|
||||
Xu, Yechen and
|
||||
Sun, Kai and
|
||||
Yu, Dian and
|
||||
Yu, Cong and
|
||||
Tian, Yin and
|
||||
Dong, Qianqian and
|
||||
Liu, Weitang and
|
||||
Shi, Bo and
|
||||
Cui, Yiming and
|
||||
Li, Junyi and
|
||||
Zeng, Jun and
|
||||
Wang, Rongzhao and
|
||||
Xie, Weijian and
|
||||
Li, Yanting and
|
||||
Patterson, Yina and
|
||||
Tian, Zuoyu and
|
||||
Zhang, Yiwen and
|
||||
Zhou, He and
|
||||
Liu, Shaoweihua and
|
||||
Zhao, Zhe and
|
||||
Zhao, Qipeng and
|
||||
Yue, Cong and
|
||||
Zhang, Xinrui and
|
||||
Yang, Zhengliang and
|
||||
Richardson, Kyle and
|
||||
Lan, Zhenzhong",
|
||||
booktitle = "Proceedings of the 28th International Conference on Computational Linguistics",
|
||||
month = dec,
|
||||
year = "2020",
|
||||
address = "Barcelona, Spain (Online)",
|
||||
publisher = "International Committee on Computational Linguistics",
|
||||
url = "https://aclanthology.org/2020.coling-main.419",
|
||||
doi = "10.18653/v1/2020.coling-main.419",
|
||||
pages = "4762--4772",
|
||||
}""",
|
||||
descriptive_stats={"n_samples": None, "avg_character_length": None},
|
||||
)
|
||||
|
||||
def dataset_transform(self):
|
||||
self.dataset = self.dataset.rename_column("sent1", "sentence1")
|
||||
self.dataset = self.dataset.rename_column("sent2", "sentence2")
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user