Update README.md

README.md

@@ -33,13 +33,10 @@ pip install torch==2.0.0
 pip install einops==0.6.1
 ```
 
-Download [h2oai_pipeline.py](h2oai_pipeline.py), store it alongside your notebook, and construct the pipeline yourself from the loaded model and tokenizer:
-
-
 ```python
 import torch
-from h2oai_pipeline import H2OTextGenerationPipeline
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
 
 tokenizer = AutoTokenizer.from_pretrained(
     "h2oai/h2ogpt-gm-oasst1-multilang-2048-falcon-7b",
@@ -47,13 +44,15 @@ tokenizer = AutoTokenizer.from_pretrained(
     padding_side="left",
     trust_remote_code=True,
 )
-model = AutoModelForCausalLM.from_pretrained(
-    "h2oai/h2ogpt-gm-oasst1-multilang-2048-falcon-7b",
-    torch_dtype=torch.bfloat16,
-    device_map={"": "cuda:0"},
+
+generate_text = pipeline(
+    model="h2oai/h2ogpt-gm-oasst1-multilang-2048-falcon-7b",
+    tokenizer=tokenizer,
+    torch_dtype=torch.float16,
     trust_remote_code=True,
+    use_fast=False,
+    device_map={"": "cuda:0"},
 )
-generate_text = H2OTextGenerationPipeline(model=model, tokenizer=tokenizer)
 
 res = generate_text(
     "Why is drinking water so healthy?",
@@ -78,6 +77,41 @@ print(generate_text.preprocess("Why is drinking water so healthy?")["prompt_text
 <|prompt|>Why is drinking water so healthy?<|endoftext|><|answer|>
 ```
 
+Alternatively, you can download [h2oai_pipeline.py](h2oai_pipeline.py), store it alongside your notebook, and construct the pipeline yourself from the loaded model and tokenizer:
+
+
+```python
+import torch
+from h2oai_pipeline import H2OTextGenerationPipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained(
+    "h2oai/h2ogpt-gm-oasst1-multilang-2048-falcon-7b",
+    use_fast=False,
+    padding_side="left",
+    trust_remote_code=True,
+)
+model = AutoModelForCausalLM.from_pretrained(
+    "h2oai/h2ogpt-gm-oasst1-multilang-2048-falcon-7b",
+    torch_dtype=torch.bfloat16,
+    device_map={"": "cuda:0"},
+    trust_remote_code=True,
+)
+generate_text = H2OTextGenerationPipeline(model=model, tokenizer=tokenizer)
+
+res = generate_text(
+    "Why is drinking water so healthy?",
+    min_new_tokens=2,
+    max_new_tokens=1024,
+    do_sample=False,
+    num_beams=1,
+    temperature=float(0.3),
+    repetition_penalty=float(1.2),
+    renormalize_logits=True
+)
+print(res[0]["generated_text"])
+```
+
 You may also construct the pipeline from the loaded model and tokenizer yourself and consider the preprocessing steps:
 
 ```python