File size: 4,101 Bytes
513825d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
# 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)
``` |