Update README.md

README.md

@@ -38,40 +38,78 @@ dtype: bfloat16
 
 ## 💻 Usage
 
-/still troubleshooting the jais family tokenizer deployment with trust remote code, in this merge/
+/Due to the jais family tokenizer deployment with trust remote code, especially if handling Arabic, the following implementation is suggested for inferencing this merge model/
 
 ```python
 !pip install -qU transformers accelerate
 
-from transformers import AutoTokenizer
-import transformers
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 
-model = "Solshine/Jais-590m-merged"
-messages = [{"role": "user", "content": "What is a large language model?"}]
+# Model and message setup
+model_name = "Solshine/Jais-590m-merged"
+user_message = "Explain how transformers work in machine learning"  # This can be any user input
 
-tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
+# Structure the message with role-content pairing for compatibility with Jais-chat format
+messages = [{"role": "user", "content": user_message}]
 
-# Manually apply a basic chat template since it's not provided by the model
+# Initialize tokenizer with trust_remote_code for custom Arabic-English handling
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+
+# Check if tokenizer is valid
+if tokenizer is None:
+    raise ValueError("Tokenizer initialization failed!")
+
+# Custom chat template including assistant role
 def custom_chat_template(messages):
     chat_prompt = ""
     for message in messages:
         role = message["role"]
         content = message["content"]
         chat_prompt += f"{role}: {content}\n"
+    # Add assistant role to prompt the model's response
+    chat_prompt += "assistant:"
     return chat_prompt
 
+# Generate the prompt
 prompt = custom_chat_template(messages)
+print(f"Generated prompt:\n{prompt}")
+
+# Initialize the model
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+if model is None:
+    raise ValueError("Model initialization failed!")
 
-# Trust Remote Code to use the unique tokenizer in the jais family series of models
-pipeline = transformers.pipeline(
+# Move model to the appropriate device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+# Initialize the text generation pipeline
+text_gen_pipeline = pipeline(
     "text-generation",
     model=model,
+    tokenizer=tokenizer,
+    device=device,
     torch_dtype=torch.float16,
-    device_map="auto",
-    trust_remote_code=True,
+    trust_remote_code=True
 )
 
-outputs = pipeline(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
-print(outputs[0]["generated_text"])
+# Generate text
+try:
+    outputs = text_gen_pipeline(
+        prompt,
+        max_new_tokens=256,
+        do_sample=True,
+        temperature=0.7,
+        top_k=50,
+        top_p=0.95,
+        pad_token_id=tokenizer.eos_token_id  # Ensure proper stopping
+    )
+    # Extract and print the assistant's response
+    generated_text = outputs[0]["generated_text"]
+    assistant_response = generated_text.split("assistant:")[1].strip()
+    print(f"Assistant's response:\n{assistant_response}")
+except Exception as e:
+    print(f"Error during text generation: {e}")
+
 ```