Update README.md

README.md

@@ -49,32 +49,29 @@ Prerequisites:
 - clone this repo. "cd" to the cloned repo.
 ```python
 >>> import os
->>> import torch
 >>> from mltools import dmx
->>> from transformers import pipeline,AutoModelForCausalLM
+>>> from transformers import pipeline
 >>> import evaluate
 >>> from datasets import load_dataset
 
-# Get model
->>> my_hf_token = os.environ.get("Dmatrix_HF_Token")
->>> device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+>>> my_hf_token = os.environ.get("HUGGING_FACE_HUB_TOKEN")
 
 >>> pipe = pipeline(
 >>>     "text-generation",
 >>>     model="d-matrix/distilgpt2",
->>>     device=device,
 >>>     use_auth_token=my_hf_token,
+>>>     trust_remote_code=True,
+>>>     # device_map="auto", # use this line for enabling pipeline parallel
+>>> )
+>>> pipe.model = dmx.Model(
+>>>     pipe.model, monkey_patched=False, hf=True, input_names=["input_ids", "labels"]
 >>> )
->>> pipe.model = dmx.Model(pipe.model,monkey_patched=False,hf=True,input_names=["input_ids", "labels"])
 
-# Configure quantization formats
->>> pipe.model.transform('FALLBACK.yaml')
+>>> pipe.model.transform("FALLBACK.yaml")
 
-# Evaluate
 >>> perplexity = evaluate.load("d-matrix/dmx_perplexity", module_type="metric")
 >>> input_texts = load_dataset("ptb_text_only", "penn_treebank", split="test")["sentence"]
->>> pipe.model.eval()
->>> results = perplexity.compute(model=pipe.model.body,references=input_texts)
+>>> results = perplexity.compute(model=pipe.model.body, references=input_texts)
 >>> print(results)
 {'loss': 4.164604187011719, 'perplexity': 64.36720275878906}
 ```