Update README.md

README.md

@@ -11,7 +11,7 @@ pipeline_tag: text2text-generation
 This Natural Language Processing (NLP) model is made available under the Apache License, Version 2.0. You are free to use, modify, and distribute this software according to the terms and conditions of the Apache 2.0 License. For the full license text, please refer to the Apache 2.0 License.
 # Usage and Specific Capabilities
 ## Text Length Limitation
-The model is optimized to analyze texts containing up to 2048 tokens. If your text exceeds this limit, we recommend splitting it into smaller chunks, each containing no more than 512 tokens. Each chunk can then be processed separately.
+The model is optimized to analyze texts containing up to 2048 tokens. If your text exceeds this limit, we recommend splitting it into smaller chunks, each containing no more than 2048 tokens. Each chunk can then be processed separately.
 ## Supported Languages
 Bulgarian, Chinese, Czech, Dutch, English, Estonian, Finnish, French, German, Greek, Indonesian, Italian, Japanese, Korean, Lithuanian, Norwegian, Polish, Portuguese, Romanian, Russian, Slovak, Spanish, Swedish, Turkish
 
@@ -26,18 +26,18 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 tokenizer = AutoTokenizer.from_pretrained("metricspace/EntityAnonymization-3B-V0.9")
 model = AutoModelForCausalLM.from_pretrained("metricspace/EntityAnonymization-3B-V0.9", torch_dtype=torch.bfloat16)
 
-def extract_assistant_response(input_text):
-    # Find all occurrences of "ASSISTANT:" in the input text
-    matches = re.finditer(r'ASSISTANT:', input_text)
+import re
+
+def extract_last_assistant_response(input_text):
+    # Find the occurrence of "ASSISTANT:" in the input text
+    match = re.search(r'ASSISTANT:', input_text)
+
+    # Get the index where the last "ASSISTANT:" ends
+    start_index = match.end()
+    response = input_text[start_index:].strip()
+    return response
 
-    # Extract text after each occurrence of "ASSISTANT:"
-    assistant_responses = []
-    for match in matches:
-        start_index = match.end()  # Get the index where "ASSISTANT:" ends
-        response = input_text[start_index:].strip()
-        assistant_responses.append(response)
 
-    return assistant_responses
 
 
 text_to_anonymize = "Sophia had always been enchanted by Venice, a historic city nestled in the heart of the Venetian lagoon. She had explored Venice on numerous occasions, each visit revealing hidden treasures in the enchanting city. On her latest trip, Sophia met Marco, a local historian, who shared captivating stories about the history of Venice.""
@@ -49,7 +49,7 @@ output_entities = model.generate(inputs.input_ids, max_new_tokens=250, do_sample
 output_entities_text = tokenizer.decode(output_entities[0], skip_special_tokens=True)
 
 # extracting entities text from assistant response
-generated_part = extract_assistant_response(output_text_1)[0]
+generated_part = extract_assistant_response(output_text_1)
 
 prompt_2 = f"USER: Rephrase with {generated_part}: {text_to_anonymize}\n\nASSISTANT:"
 inputs = tokenizer(prompt_2, return_tensors='pt').to('cuda')