|
--- |
|
|
|
|
|
{} |
|
--- |
|
|
|
# Relation Extraction model for KBQA |
|
|
|
<!-- Provide a quick summary of what the model is/does. --> |
|
|
|
This is the fine-tuned version of TinyLlama for the Relation Extraction task. |
|
|
|
## Model Details |
|
The model is trained on the [Dataset](https://github.com/dki-lab/GrailQA/tree/main/data), which consists of questions and their related information, such as entities and relations. The relationships in the data are annotated from the Freebase dataset. |
|
The relationships in the data are annotated from the Freebase dataset. |
|
|
|
#### How to use |
|
You will need the transformers>=4.34 |
|
Do check the [TinyLlama](https://github.com/jzhang38/TinyLlama) github page for more information. |
|
|
|
|
|
### Direct Use |
|
|
|
```python |
|
import os |
|
import torch |
|
#from datasets import load_dataset |
|
from transformers import ( |
|
AutoModelForCausalLM, |
|
AutoTokenizer, |
|
BitsAndBytesConfig, |
|
HfArgumentParser, |
|
TrainingArguments, |
|
pipeline, |
|
logging, |
|
) |
|
|
|
from transformers import AutoTokenizer, pipeline, logging |
|
model_name_or_path = "dice-research/Ft_TinnyLlama_QA_RE" |
|
model_basename = "model" |
|
|
|
use_triton = False |
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) |
|
tokenizer.pad_token = tokenizer.eos_token |
|
model = AutoModelForCausalLM.from_pretrained( |
|
model_name_or_path, quantization_config=None, device_map="auto" |
|
) |
|
|
|
pipe = pipeline("text-generation", model=model,tokenizer=tokenizer, torch_dtype=torch.bfloat16, device_map="auto") |
|
|
|
def prompt_REQA(question): |
|
messages = [ |
|
{"role": "user", "content": question}, |
|
] |
|
|
|
prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
|
#response=pipe(sparql_prompt) |
|
response = pipe(prompt,max_new_tokens=20, do_sample=True, temperature=0.6, top_k=5, top_p=0.95)[0]['generated_text'] |
|
return response.split('<|assistant|>\n')[1] |
|
|
|
prompt_REQA("how many electronic arts games are available for sale in the united states of america?") |
|
``` |
|
|