Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import torch
+import re
 import os
 auth_token = os.environ.get("HUGGING_FACE_HUB_TOKEN")
 from unsloth import FastLanguageModel
@@ -29,6 +30,26 @@ From the given email, extract the following key values. The keys are explained b
 """
 
 
+
+def process_output(output):
+    """
+    Process the output to extract the response.
+    """
+    # Define the regex pattern
+    pattern = r'### Response:\n?(.*?)<\|endoftext\|>'
+    # Search for the pattern in the output
+    match = re.search(pattern, output, re.DOTALL)
+    
+    if match:
+        # Extract the response
+        response = match.group(1)
+        # Remove specified symbols
+        cleaned_str = re.sub(r'\\n|\\\\|\\\'', '', response)
+        return cleaned_str
+    else:
+        return output
+
+
 # Define the function for generating output based on input
 def generate_output(input_text,model):
     # Prompt for the instruction
@@ -59,7 +80,8 @@ def generate_output(input_text,model):
     # Generate outputs
     outputs = model.generate(**inputs, max_new_tokens=2048, use_cache=True)
     output = tokenizer.batch_decode(outputs)
-    return output
+    cleaned_response = process_output(output)
+    return output,cleaned_response
 
 
 model_options = ["sxandie/llama_3_8b_4bitQ","DataIntelligenceTeam/NER-Phi-3-mini-4k-instruct"]
@@ -68,10 +90,15 @@ inputs = [
     gr.inputs.Dropdown(label="Model", choices=model_options, default=model_options[0])
 ]
 
+outputs = [
+    gr.outputs.Textbox(label="Original Output Text"),
+    gr.outputs.Textbox(label="Formatted JSON")
+]
+
 # Create Gradio interface
 iface = gr.Interface(fn=generate_output, 
                      inputs=inputs, 
-                     outputs="text", 
+                     outputs= outputs,
                      title="Email Information Extraction", 
                      description="Extract key information from the provided email.")
 iface.launch()