|
# Ahma-3B-RAG |
|
|
|
## Overview |
|
Ahma-3B-RAG is a 3B-parameter language model fine-tuned on **Retrieval-Augmented Generation (RAG) problems** using approximately **20,000 synthetically generated samples**. The synthetic data was created using **Nemotron-70B** and **DeepSeekV3** to improve the model's ability to handle RAG-based tasks effectively. |
|
|
|
## Model Information |
|
- **Model Name:** Ahma-3B-RAG |
|
- **Training Data:** ~20k synthetic RAG samples (Nemotron-70B, DeepSeekV3) |
|
- **Use Case:** RAG-based response generation |
|
- **Primary Language:** Finnish |
|
|
|
## Installation & Dependencies |
|
Before using the model, make sure you have the necessary dependencies installed: |
|
|
|
```bash |
|
pip install torch transformers |
|
``` |
|
|
|
```python |
|
# Tests were run with the following package versions |
|
# You can try with different versions as well but these should at least work |
|
import transformers |
|
import flash_attn |
|
import torch |
|
|
|
assert transformers.__version__ == 4.48.1 |
|
assert torch.__version__ == 2.1.2+cu121 |
|
assert flash_attn.__version__ == 2.7.3 |
|
``` |
|
|
|
## Model Loading |
|
To load the model efficiently, use the following function: |
|
|
|
```python |
|
import torch |
|
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig |
|
|
|
def load_llama_model(model_path, max_seq_length=2048, dtype=None): |
|
""" |
|
Loads the LLaMA model with the given configuration. |
|
|
|
Args: |
|
model_path (str): Path or name of the pre-trained model. |
|
max_seq_length (int): Maximum sequence length for the model. |
|
dtype (torch.dtype or None): Data type for the model. Default is auto-detected. |
|
|
|
Returns: |
|
model, tokenizer, generation_config: Loaded model, tokenizer, and generation config. |
|
""" |
|
# Set default dtype based on available hardware |
|
torch_dtype = torch.bfloat16 if dtype is None else dtype |
|
|
|
# Load model with appropriate configuration |
|
model = AutoModelForCausalLM.from_pretrained( |
|
model_path, |
|
torch_dtype=torch_dtype, |
|
device_map='auto', |
|
attn_implementation="flash_attention_2" # If you do not have access to GPU supporting flash_attention_2 you can commit this line |
|
) |
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path) |
|
|
|
generation_config = GenerationConfig( |
|
pad_token_id=tokenizer.eos_token_id, |
|
eos_token_id=tokenizer.convert_tokens_to_ids("</s>") |
|
) |
|
|
|
return model, tokenizer, generation_config |
|
|
|
model_path = "RASMUS/AHMA-3B-RAG" |
|
``` |
|
|
|
## Generating Prompts for RAG |
|
To generate prompts that incorporate context for RAG-based queries, use the following function: |
|
|
|
```python |
|
def generate_rag_prompt_message(row): |
|
prompt = f'Olet tekoälyavustaja joka vastaa annetun kontekstin perusteella asiantuntevasti ja ystävällisesti käyttäjän kysymyksiin\n\nKonteksti: {row["text"]}\n\nKysymys: {row["question"]}\n\nVastaa yllä olevaan kysymykseen annetun kontekstin perusteella.' |
|
row["messages"] = [{'role': 'user', 'content': prompt}] |
|
return row |
|
``` |
|
|
|
## Generating Responses |
|
Ahma-3B-RAG can be used to generate responses using the following inference setup: |
|
|
|
```python |
|
model, tokenizer, generation_config = load_llama_model(model_path) |
|
|
|
row = {"text": "Rasmus Toivanen loi tämän mallin", "question": "Kuka loi tämän mallin?"} |
|
row = generate_rag_prompt_message(row) |
|
|
|
inputs = tokenizer( |
|
[ |
|
tokenizer.apply_chat_template(row["messages"], tokenize=False) |
|
] * 1, return_tensors="pt" |
|
).to("cuda") |
|
|
|
with torch.no_grad(): |
|
generated_ids = model.generate( |
|
input_ids=inputs["input_ids"], |
|
attention_mask=inputs["attention_mask"], |
|
generation_config=generation_config, **{ |
|
"temperature": 0.1, |
|
"penalty_alpha": 0.6, |
|
"min_p": 0.3, |
|
"do_sample": True, |
|
"max_new_tokens": 300 |
|
} |
|
) |
|
|
|
generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True)[0] |
|
generated_text_cleaned = generated_text.split('[/INST]')[1].replace('</s>', '').strip() if '[/INST]' in generated_text else generated_text.strip() |
|
|
|
print(generated_text_cleaned) |
|
``` |