Update README.md and config.

README.md

@@ -12,17 +12,30 @@ We provide a joint model card for `Pharia-1-LLM-7B-control` and `Pharia-1-LLM-co
 # Usage
 
 ```python
+import torch
+
 from transformers import AutoModelForCausalLM, PreTrainedTokenizerFast
 
-INPUT = "Hello, how are you"
-MODEL_ID = "Aleph-Alpha/Pharia-1-LLM-7B-control-safetensors"
+
+INPUT = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+You are a helpful assistant. You give engaging, well-structured answers to user inquiries.<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+When was Rome founded?<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+
+
+"""
+
+MODEL_ID = "Aleph-Alpha/Pharia-1-LLM-7B-control-hf"
 
 tokenizer = PreTrainedTokenizerFast.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True, torch_dtype=torch.bfloat16)
 
-inputs = tokenizer(INPUT, return_token_type_ids=False, return_tensors="pt")
-outputs = model.generate(**inputs, max_new_tokens=50)
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
 
-generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+inputs = tokenizer(INPUT, return_token_type_ids=False, return_tensors="pt").to(device)
+outputs = model.generate(**inputs, max_new_tokens=50)
+generated_text = tokenizer.decode(outputs[0])
 print(generated_text)
 ```

config.json

@@ -24,7 +24,7 @@
   "rope_scaling": null,
   "rope_theta": 1000000,
   "tie_word_embeddings": false,
-  "torch_dtype": "float32",
+  "torch_dtype": "bfloat16",
   "transformers_version": "4.44.2",
   "use_cache": true,
   "vocab_size": 128000

modeling_pharia.py

@@ -764,7 +764,6 @@ class PhariaForCausalLM(PhariaPreTrainedModel):
 
         hidden_states = outputs[0]
         logits = self.lm_head(hidden_states)
-        logits = logits.float()
 
         return CausalLMOutputWithPast(
             loss=0.0,