Fix inference code in readme

README.md

@@ -33,118 +33,58 @@ This repo contains a low-rank adapter for LLaMA-7b fit on the Stanford Alpaca da
 
 Model can be easily loaded with AutoModelForCausalLM.
 ``` python
-# import torch
-from peft import PeftModel
-# from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig
-
-
 import torch
-# from peft import PeftModel
+from peft import PeftModel
 import transformers
-import gradio as gr
-
-assert (
-    "LlamaTokenizer" in transformers._import_structure["models.llama"]
-), "LLaMA is now in HuggingFace's main branch.\nPlease reinstall it: pip uninstall transformers && pip install git+https://github.com/huggingface/transformers.git"
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from peft import PeftModel, PeftConfig
 from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig
 
-tokenizer = LlamaTokenizer.from_pretrained("decapoda-research/llama-7b-hf")
+base_model_path = "meta-llama/Llama-2-7b-hf"
+adapter_path = "OdiaGenAI/odiagenAI-model-v1"
 
-BASE_MODEL = "decapoda-research/llama-7b-hf"
-LORA_WEIGHTS = "OdiaGenAI/odiagenAI-model-v1"
+tokenizer = AutoTokenizer.from_pretrained(base_model_path, trust_remote_code=True)
+tokenizer.pad_token = tokenizer.eos_token
 
-if torch.cuda.is_available():
-    device = "cuda"
-else:
-    device = "cpu"
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.float16,
+)
 
-try:
-    if torch.backends.mps.is_available():
-        device = "mps"
-except:
-    pass
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_path,
+    quantization_config=bnb_config,
+    device_map="auto",
+    trust_remote_code=True
+)
 
-if device == "cuda":
-    model = LlamaForCausalLM.from_pretrained(
-        BASE_MODEL,
-        load_in_8bit=False,
-        torch_dtype=torch.float16,
-        device_map="auto",
-    )
-    model = PeftModel.from_pretrained(
-        model, LORA_WEIGHTS, torch_dtype=torch.float16, force_download=True
-    )
-elif device == "mps":
-    model = LlamaForCausalLM.from_pretrained(
-        BASE_MODEL,
-        device_map={"": device},
-        torch_dtype=torch.float16,
-    )
-    model = PeftModel.from_pretrained(
-        model,
-        LORA_WEIGHTS,
-        device_map={"": device},
-        torch_dtype=torch.float16,
-    )
-else:
-    model = LlamaForCausalLM.from_pretrained(
-        BASE_MODEL, device_map={"": device}, low_cpu_mem_usage=True
-    )
-    model = PeftModel.from_pretrained(
-        model,
-        LORA_WEIGHTS,
-        device_map={"": device},
-    )
+model = PeftModel.from_pretrained(base_model, adapter_path)
+
+instruction = "ଭାରତ ବିଷୟରେ କିଛି କୁହନ୍ତୁ"
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def generate_prompt(instruction, input=None):
-    if input:
-        return f"""### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n"""
-    else:
-        return f"""### Instruction:\n{instruction}\n\n### Response:\n"""
-
-if device != "cpu":
-    model.half()
-model.eval()
-if torch.__version__ >= "2":
-    model = torch.compile(model)
-
-def evaluate(
-    instruction,
-    input=None,
+inputs = tokenizer(instruction, return_tensors="pt").to(device)
+input_ids = inputs["input_ids"].to(device)
+generation_config = GenerationConfig(
     temperature=0.1,
     top_p=0.75,
     top_k=40,
     num_beams=4,
-    max_new_tokens=128,
-    **kwargs,
-):
-    prompt = generate_prompt(instruction, input)
-    print(prompt)
-    inputs = tokenizer(prompt, return_tensors="pt")
-    print(inputs)
-    input_ids = inputs["input_ids"].to(device)
-    print(input_ids)
-    generation_config = GenerationConfig(
-        temperature=temperature,
-        top_p=top_p,
-        top_k=top_k,
-        num_beams=num_beams,
-        **kwargs,
+)
+with torch.no_grad():
+    generation_output = model.generate(
+        input_ids=input_ids,
+        generation_config=generation_config,
+        return_dict_in_generate=True,
+        output_scores=True,
+        max_new_tokens=128,
     )
-    with torch.no_grad():
-        generation_output = model.generate(
-            input_ids=input_ids,
-            generation_config=generation_config,
-            return_dict_in_generate=True,
-            output_scores=True,
-            max_new_tokens=max_new_tokens,
-        )
-        print(generation_output)
-    s = generation_output.sequences[0]
-    print(s)
-    output = tokenizer.decode(s)
-    print(output)
-    return output.split("### Response:")[1].strip()
+s = generation_output.sequences[0]
+output = tokenizer.decode(s)
+print(output)
 
 ```