Update README.md

README.md

@@ -2,9 +2,7 @@
 tags:
 - autotrain
 - text-generation
-widget:
-- text: "I love AutoTrain because "
-license: other
+
 ---
 
 # Model Trained Using AutoTrain
@@ -15,26 +13,64 @@ This model was trained using AutoTrain. For more information, please visit [Auto
 
 ```python
 
-from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model_path = "PATH_TO_THIS_REPO"
+from peft import PeftModel, PeftConfig
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+import torch
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+config = PeftConfig.from_pretrained("Ashishkr/llama2-qrecc-context-resolution")
+model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
+model = PeftModel.from_pretrained(model, "Ashishkr/llama2-qrecc-context-resolution").to(device)
+tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
+
+def response_generate(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+    prompt: str,
+    max_new_tokens: int = 128,
+    temperature: float = 0.7,
+):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    inputs = tokenizer(
+        [prompt],
+        return_tensors="pt",
+        return_token_type_ids=False,
+    ).to(
+        device
+    )
 
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModelForCausalLM.from_pretrained(
-    model_path,
-    device_map="auto",
-    torch_dtype='auto'
-).eval()
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        response = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            return_dict_in_generate=True,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id,
+        )
 
-# Prompt content: "hi"
-messages = [
-    {"role": "user", "content": "hi"}
-]
+    decoded_output = tokenizer.decode(
+        response["sequences"][0],
+        skip_special_tokens=True,
+    )
 
-input_ids = tokenizer.apply_chat_template(conversation=messages, tokenize=True, add_generation_prompt=True, return_tensors='pt')
-output_ids = model.generate(input_ids.to('cuda'))
-response = tokenizer.decode(output_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
+    return decoded_output
+
+prompt = """>>CONTEXT<<I heard John Marks was the first christian missionary in Ireland. What was the capital then??>>REWRITE<<"""
+
+response = response_generate(
+    model,
+    tokenizer,
+    prompt,
+    max_new_tokens=20,
+    temperature=0.1,
+)
 
-# Model response: "Hello! How can I assist you today?"
 print(response)
+
+
 ```