Update README.md

README.md

@@ -1,34 +1,33 @@
 ---
 license: other
+pipeline_tag: text-generation
+tags:
+- sharded
+- yi
+- bf16
+- colab
 ---
 
+# Yi-6B-200K-Llama - sharded bf16
 
-https://huggingface.co/chargoddard/Yi-6B-200K-Llama
+https://huggingface.co/chargoddard/Yi-6B-200K-Llama sharded & in bf16 (colab free GPU loadable)
+
+Please note the custom license from [the original](https://huggingface.co/01-ai/Yi-6B-200K) still applies 🌝
 
 
 ```python
-# !pip install -U -q transformers datasets accelerate sentencepiece bitsandbytes
+# !pip install -U -q transformers accelerate sentencepiece bitsandbytes
 import torch
-from huggingface_hub import notebook_login
 
 # Load model directly
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-notebook_login()
-
-
-tokenizer = AutoTokenizer.from_pretrained("pszemraj/Yi-6B-200K-Llama-sharded")
+tokenizer = AutoTokenizer.from_pretrained("pszemraj/Yi-6B-200K-Llama-sharded", use_fast=False)
 model = AutoModelForCausalLM.from_pretrained(
     "pszemraj/Yi-6B-200K-Llama-sharded", load_in_4bit=True
 )
 
-prompt = """God, I just love smoking weed-"
-"I know, I love smoking weed too," God replied, "but what would you say if I were to tell you that the earth is 10,000 years old, that men and dinosaurs walked the earth together, that there was once a great flood, and that a man named Jonah once walked in the belly of a fish?"
-"That's some weird shit," the man said.
-"And that's the Gospel, the Good News," God continued, "that Jesus was raised from the dead, that he died for your sins, and that if you have faith in him you will go to heaven when you die."
-"What do you mean?" the man asked.
-"I mean," God answered, "that if you have any sense at all, you'll roll another joint."
-"I don't know why but I'm feeling"""
+prompt = "Custom non-commercial licenses are just so fun, aren't they? The best thing about them is"
 
 
 if torch.cuda.is_available():
@@ -46,4 +45,4 @@ tokens = model.generate(
     repetition_penalty=1.05,
 )
 print(tokenizer.decode(tokens[0], skip_special_tokens=True))
-```
+```