Update README.md

README.md

@@ -29,6 +29,11 @@ Use the code below to get started with the model.
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 
+try:
+    if torch.backends.mps.is_available():
+        device = "mps"
+except:
+    pass
 checkpoint = "abdoelsayed/llama-7b-v2-Receipt-Key-Extraction"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -37,7 +42,7 @@ tokenizer = AutoTokenizer.from_pretrained(checkpoint, model_max_length=512,
         use_fast=False,)
 model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
 
-def generate_response(instruction, input_text, max_new_tokens=100, temperature=0.1,  num_beams=4 ,top_k=40):
+def generate_response(instruction, input_text, max_new_tokens=100, temperature=0.1,  num_beams=4 , top_p=0.75, top_k=40):
     prompt = f"Below is an instruction that describes a task, paired with an input that provides further context.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n\n### Response:"
     inputs = tokenizer(prompt, return_tensors="pt")
     input_ids = inputs["input_ids"].to(device)
@@ -48,9 +53,9 @@ def generate_response(instruction, input_text, max_new_tokens=100, temperature=0
             num_beams=num_beams,
         )
     with torch.no_grad():
-        outputs = model.generate(input_ids,generation_config=generation_config, max_new_tokens=max_new_tokens)
+        outputs = model.generate(input_ids,generation_config=generation_config, max_new_tokens=max_new_tokens,return_dict_in_generate=True,output_scores=True,)
     outputs = tokenizer.decode(outputs.sequences[0])
-    return output.split("### Response:")[-1].strip().replace("</s>","")
+    return outputs.split("### Response:")[-1].strip().replace("</s>","")
 
 instruction = "Extract the class, Brand, Weight, Number of units, Size of units, Price, T.Price, Pack, Unit from the following sentence"
 input_text = "Americana Okra zero 400 gm"
@@ -58,6 +63,7 @@ input_text = "Americana Okra zero 400 gm"
 response = generate_response(instruction, input_text)
 print(response)
 
+
 ```