Update app.py

app.py

@@ -47,41 +47,28 @@ instruction = ("Instruction: As a helpful, respectful and trustworthy debiasing
 
 
 def get_debiased_sequence(prompt):
-    """Generate a debiased version of the provided text using the debiasing pipeline."""
-    instruction_prefix = "<s> <<SYS>> {instruction} <</SYS>> [INST]".format(instruction=instruction)
-    instruction_suffix = "[/INST]</s>"
-    full_input_text = f"{instruction_prefix}{prompt}{instruction_suffix}"
-
-    # Tokenize the full input text to calculate its length in tokens
-    input_tokens = debias_tokenizer.encode(full_input_text)
-
-    # Ensure max_length is greater than the number of input tokens
-    max_length = len(input_tokens) + 50  # Add a buffer to accommodate generation without truncation
-
-    try:
-        sequences = debias_pipeline(
-            full_input_text,
-            do_sample=True,
-            top_k=10,
-            num_return_sequences=1,
-            eos_token_id=debias_tokenizer.eos_token_id,
-            max_length=max_length,  # Updated to use calculated max_length
-        )
-
-        if sequences:
-            res = sequences[0]['generated_text']
-            # Assuming the response also includes the [/INST] tag, split and extract after this tag
-            result_part = res.split('[/INST]')[-1]
-            clean_result = ''.join(c for c in result_part if c.isprintable())
-            return clean_result.strip()
-    except RuntimeError as e:
-        if 'CUDA out of memory' in str(e):
-            torch.cuda.empty_cache()  # Try clearing cache to free up memory
-            return "Error: Out of memory. Please try again with shorter input or less complex instructions."
-        else:
-            raise e  # Re-raise the exception if it's not a memory error
-
-    return "No output generated. Check model configuration or input."
+    """
+    Generate a debiased version of the provided text using the debiasing pipeline.
+    Args:
+    - prompt (str): Text to be debiased.
+    Returns:
+    - str: Debiased text.
+    """
+    input_text = f"<s> <<SYS>>{instruction}. {sys_message} <</SYS>> [INST]{prompt} [/INST]"
+    sequences = debias_pipeline(
+        input_text,
+        do_sample=True,
+        top_k=10,
+        num_return_sequences=1,
+        eos_token_id=tokenizer.eos_token_id,
+        max_length=len(prompt.split(" ")) + len(input_text.split(" ")) + 100,  # Increased max_length
+    )
+    res = sequences[0]['generated_text']
+    result_part = res.split('[/INST]')[-1]
+    clean_result = ''.join(c for c in result_part if c.isprintable())
+    cleaned_text = re_incomplete_sentence(clean_result.strip())
+    return cleaned_text.strip()
+
 
 # Streamlit interface setup
 st.title('UnBIAS App')