Update README.md

README.md

@@ -18,7 +18,7 @@ from peft import PeftModel, PeftConfig
 from transformers import AutoModelForCausalLM
 from transformers import AutoTokenizer
 import torch
-
+import re
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 config = PeftConfig.from_pretrained("Ashishkr/llama2-qrecc-context-resolution")
@@ -70,7 +70,11 @@ response = response_generate(
     temperature=0.1,
 )
 
-print(response)
+def extract_between_tags(input_string):
+    pattern = r'>>REWRITE<<(.*?)</REWRITE>'
+    match = re.search(pattern, input_string)
+    return match.group(1) if match else ''
 
+print(extract_between_tags(response))
 
 ```