Update README.md

README.md

@@ -1 +1,115 @@
-Instruction and Testing results coming later...
+# 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)
+```