Update README

README.md

@@ -39,22 +39,22 @@ LIME-1B is a 1B-parameter, decoder-only Transformer language model trained from
 
 LIME-1B follows a modern GPT-style decoder-only Transformer with several quality-oriented design choices:
 
-| Component                   | Value                |
-|-----------------------------|----------------------|
-| Architecture                | Decoder-only Transformer |
-| Parameters                  | 1.0B                 |
-| Layers (decoder blocks)     | 32                   |
-| d_model                     | 1536                 |
-| FFN dimension (d_ff)        | 6144                 |
-| Attention heads             | 24                   |
-| Vocabulary size             | 50,000               |
-| Max sequence length         | 512 tokens           |
-| Positional encoding         | Sinusoidal           |
-| Norm                        | RMSNorm              |
-| FFN                         | SiLU MLP             |
-| Attention                   | FlashAttention   |
-| Tying of embeddings         | Output head tied to embedding |
-| Precision (training)        | Mixed fp32/bf16 (autocast) + grad clipping |
+| Component               | Value                                      |
+|-------------------------|--------------------------------------------|
+| Architecture            | Decoder-only Transformer                   |
+| Parameters              | 1.0B                                       |
+| Layers (decoder blocks) | 32                                         |
+| d_model                 | 1536                                       |
+| FFN dimension (d_ff)    | 6144                                       |
+| Attention heads         | 24                                         |
+| Vocabulary size         | 50,000                                     |
+| Max sequence length     | 512 tokens                                 |
+| Positional encoding     | Sinusoidal                                 |
+| Norm                    | RMSNorm                                    |
+| FFN                     | SiLU MLP                                   |
+| Attention               | FlashAttention                             |
+| Tying of embeddings     | Output head tied to embedding              |
+| Precision (training)    | Mixed fp32/bf16 (autocast) + grad clipping |
 
 
 ## 2. Training data
@@ -121,24 +121,28 @@ After pretraining, the model is fine-tuned on a **unified instruction schema**:
 
 ## Usage
 ```python
+# Example usage
+# pip install -U ukraine
+
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-model_name = "anarlavrenov/LIME-1B"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+model_name = "anarlavrenov/LIME-1b"
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype=torch.bfloat16,
     device_map="auto",
+    trust_remote_code=True
 )
 
-def build_inference_prompt(context, question):
+def build_prompt(context_, question_, tokenizer_):
 
   context_ids  = tokenizer.encode(context) if context else []
   question_ids = tokenizer.encode(question)
 
-  uid = tokenizer.convert_tokens_to_ids("<assistant>")
-  aid = tokenizer.convert_tokens_to_ids("<user>")
+  uid = tokenizer.convert_tokens_to_ids("<user>")
+  aid = tokenizer.convert_tokens_to_ids("<assistant>")
 
   ids = []
 
@@ -150,18 +154,17 @@ def build_inference_prompt(context, question):
 
   return torch.tensor(ids, dtype=torch.long)
 
-# Example usage
 context = "..."  # optional context
 question = "Write five questions for a Data Scientist interview."
-prompt = build_prompt(context, question)
+prompt = build_prompt(context, question, tokenizer)
 
 inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 outputs = model.generate(
     **inputs,
     max_new_tokens=256,
     do_sample=True,
-    top_p=0.9,
-    temperature=0.5,
+    top_p=None,
+    temperature=None,
     pad_token_id=tokenizer.pad_token_id,
     eos_token_id=tokenizer.eos_token_id,
 )