Update README.md

README.md

@@ -1,7 +1,6 @@
 ---
 tags:
 - autotrain
-- conversational
 - meta-llama
 - meta-llama/Llama-2-7b-hf
 inference: true
@@ -18,36 +17,56 @@ widget:
 
     response: ''
 library_name: peft
+pipeline_tag: text-generation
 ---
 
-```python
-!huggingface-cli login
-
- 
-    _|    _|  _|    _|    _|_|_|    _|_|_|  _|_|_|  _|      _|    _|_|_|      _|_|_|_|    _|_|      _|_|_|  _|_|_|_|
-    _|    _|  _|    _|  _|        _|          _|    _|_|    _|  _|            _|        _|    _|  _|        _|
-    _|_|_|_|  _|    _|  _|  _|_|  _|  _|_|    _|    _|  _|  _|  _|  _|_|      _|_|_|    _|_|_|_|  _|        _|_|_|
-    _|    _|  _|    _|  _|    _|  _|    _|    _|    _|    _|_|  _|    _|      _|        _|    _|  _|        _|
-    _|    _|    _|_|      _|_|_|    _|_|_|  _|_|_|  _|      _|    _|_|_|      _|        _|    _|    _|_|_|  _|_|_|_|
-    
-    To login, `huggingface_hub` requires a token generated from https://huggingface.co/settings/tokens .
-Token: <your-hf-access-token>
-```
-
 
 ```python
-
+import transformers
 from peft import PeftModel, PeftConfig
-from transformers import AutoModelForCausalLM
-from transformers import AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+from torch import cuda, bfloat16
+
+base_model_id = 'meta-llama/Llama-2-7b-chat-hf'
+
+device = f'cuda:{cuda.current_device()}' if cuda.is_available() else 'cpu'
+
+bnb_config = transformers.BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type='nf4',
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=bfloat16
+)
+
+
+hf_auth = "your-huggingface-access-token"
+model_config = transformers.AutoConfig.from_pretrained(
+    base_model_id,
+    use_auth_token=hf_auth
+)
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = transformers.AutoModelForCausalLM.from_pretrained(
+    base_model_id,
+    trust_remote_code=True,
+    config=model_config,
+    quantization_config=bnb_config,
+    device_map='auto',
+    use_auth_token=hf_auth
+)
 
 config = PeftConfig.from_pretrained("Ashishkr/llama2_medical_consultation")
-model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
 model = PeftModel.from_pretrained(model, "Ashishkr/llama2_medical_consultation").to(device)
-tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
+
+model.eval()
+print(f"Model loaded on {device}")
+
+tokenizer = transformers.AutoTokenizer.from_pretrained(
+    model_id,
+    use_auth_token=hf_auth
+)
+
+
 
 ```
 
@@ -60,7 +79,6 @@ def llama_generate(
     prompt: str,
     max_new_tokens: int = 128,
     temperature: float = 0.92):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     inputs = tokenizer(
         [prompt],
@@ -70,7 +88,10 @@ def llama_generate(
         device
     )
 
-    with torch.autocast("cuda", dtype=torch.bfloat16):
+    # Check if bfloat16 is supported, otherwise use float16
+    dtype_to_use = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+
+    with torch.autocast("cuda", dtype=dtype_to_use):
         response = model.generate(
             **inputs,
             max_new_tokens=max_new_tokens,
@@ -87,7 +108,6 @@ def llama_generate(
 
     return decoded_output[len(prompt) :]
 
-
 prompt = """
  instruction: "If you are a doctor, please answer the medical questions based on the patient's description." \n
 
@@ -97,7 +117,7 @@ goes in my left arm/hand/fingers. I have had headaches since the aneurysm,
 but this is different. Also, my moods have been horrible for the past few weeks.\n
 
 response:  """
-
+# You can use the function as before
 response = llama_generate(
     model,
     tokenizer,