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:

pip install torch transformers

# 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:

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:

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:

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)