Update README.md

README.md

@@ -143,9 +143,10 @@ input_args = {
 import torch
 import transformers
 
-model_name = "BUT-FIT/csmpt-6.7B-RAGsum"
+model_name = "BUT-FIT/csmpt-7B-RAGsum"
 tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
 config = transformers.AutoConfig.from_pretrained(model_name, trust_remote_code=True)
+config.init_device = 'cuda:0'  # For fast initialization directly on GPU!
 
 formatted_input = f"""Shrň následující výsledky pro dotaz "{input_args['query']}".
 |Výsledky|: {input_args['input']}
@@ -157,6 +158,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 model = transformers.AutoModelForCausalLM.from_pretrained(
     model_name,
     config=config,
+    torch_dtype=torch.bfloat16,  # Load model weights in bfloat16
     trust_remote_code=True
 ).cuda()
 with torch.autocast('cuda', dtype=torch.bfloat16):
@@ -173,8 +175,6 @@ with torch.autocast('cuda', dtype=torch.bfloat16):
     input_length = inputs['input_ids'].shape[1]
     generated_text = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
 
-    print("RAG Summary:", generated_text)
-
 ```
 
 Example of generated summary