Updated inference script

README.md

@@ -37,7 +37,7 @@ Krutrim Large Language Model (LLM) is a 2 trillion token multilingual foundation
 
 | Model Name | Release Date |Release Note | Reference|
 |------------|-------------|-------------|-------------|
-| Krutrim-1-Instruct  | 2024-01-31 | SFT on Krutrim-1-Base |[Here](https://huggingface.co/krutrim-ai-labs/Krutrim-1-instruct)
+| Krutrim-1-Instruct  | 2024-01-31 | SFT on Krutrim-1 Base |[Here](https://huggingface.co/krutrim-ai-labs/Krutrim-1-instruct)
 
 
 ## Data Freshness
@@ -110,42 +110,34 @@ Krutrim Large Language Model (LLM) is a 2 trillion token multilingual foundation
 
 ## Usage
 
-To run this model, do this:
-```
-git clone https://github.com/ola-krutrim/Krutrim-1-7B.git
-cd Krutrim-1-7B
-pip install -r requirements.txt
-```
-
-To test the base model, you can run
-```
-python inference/inference.py
-```
-
-To test batch inference of instruct model, you can run
-```
-python inference/batch_inference.py
-```
-
 To use the instruct model, you can load it with `AutoModelForCausalLM` as follows:
 ```
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model_id = "krutrim-ai-labs/Krutrim-1-base"
+model_id = "krutrim-ai-labs/Krutrim-1-instruct"
 # Load model and tokenizer
 model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, trust_remote_code=True)
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-prompt = "Hello"
+# Apply Chat Template
+chat_template ="{% for message in messages %}{% if message['role'] == 'system' %}{{ '<|SYSTEM|> ' + message['content'] + '\n' }}{% elif message['role'] == 'user' %}{{ '<|USER|> ' + message['content'] + '\n' }}{% elif message['role'] == 'assistant' %}{% if not loop.last %}{{ '<|RESPONSE|>\n'  + message['content'] + eos_token + '\n' }}{% else %}{{ '<|RESPONSE|>\n'  + message['content'] + eos_token }}{% endif %}{% endif %}{% if loop.last and add_generation_prompt %}{{ '<|RESPONSE|>\n' }}{% endif %}{% endfor %}"
+tokenizer.chat_template = chat_template
+
+prompt_dict = [
+
+    {"role": "system", "content": "You are an AI assistant."},
+    {"role": "user", "content": "Who are you?"}
+]
 
-inputs = tokenizer(prompt, return_tensors='pt')
+prompts = tokenizer.apply_chat_template(prompt_dict, add_generation_prompt=True, tokenize=False)
+inputs = tokenizer(prompts, return_tensors='pt').to(device)
 inputs.pop("token_type_ids", None)
 
 # Generate response
 outputs = model.generate(
     **inputs,
-    max_length=5
+    max_length=100
 )
 
 response = tokenizer.decode(outputs[0])