metadata
license: apache-2.0
language:
- en
tags:
- code
- knowledge extraction
- tiny
- small
A model that can extract the knowledge points from the given C language code.
The base model is pythia-70m. This model was fine-tuned with 10 epochs using Q-Lora method on my own training set.
A usage example is as follows, first import the model and prepare the code:
from transformers import GPTNeoXForCausalLM, AutoTokenizer
model_name_or_path = 'Mxode/Pythia-70m-C-Language-KnowledgeExtract'
device = 'cuda'
model = GPTNeoXForCausalLM.from_pretrained(model_name_or_path).to(device)
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
instruction = '[Summarize the knowledge points in the code below]\n' # instruction template
# any c-lang pieces you like, could be partial functions or statements
input_content = '''```c
int partition(int arr[], int low, int high) {
int pivot = arr[high];
int i = (low - 1);
for (int j = low; j <= high - 1; j++) {
if (arr[j] < pivot) {
i++;
swap(&arr[i], &arr[j]);
}
}
swap(&arr[i + 1], &arr[high]);
return (i + 1);
}
void quickSort(int arr[], int low, int high) {
if (low < high) {
int pi = partition(arr, low, high);
quickSort(arr, low, pi - 1);
quickSort(arr, pi + 1, high);
}
}
```'''
text = instruction + input_content
Then generate:
inputs = tokenizer(text, return_tensors="pt").to(device)
tokens = model.generate(
**inputs,
pad_token_id=tokenizer.eos_token_id,
max_new_tokens=32,
)
response = tokenizer.decode(tokens[0]).split('```')[-1].split('<')[0] # deduplicate inputs
However, in practical use, in order to achieve more diverse representations, it's recommended to do multiple inferences. Don't worry, it's really small so the inferences don't take much time, as follows:
ans_dict = {}
def increment_insert(key):
ans_dict[key] = ans_dict.get(key, 0) + 1
for i in range(30): # maybe 20 times or less enough too
inputs = tokenizer(text, return_tensors="pt").to(device)
tokens = model.generate(
**inputs,
pad_token_id=tokenizer.eos_token_id,
max_new_tokens=32,
do_sample=True,
temperature=2.0, # high temperature for diversity
top_p=0.95,
top_k=30,
)
response = tokenizer.decode(tokens[0]).split('```')[-1].split('<')[0]
increment_insert(response)
print(ans_dict)
### output as below, could take high-freq answers
### {'Backtracking': 1, 'Heap': 1, 'Quick sort': 25, 'Recurrence': 2, 'Queue': 1}