Update app.py

app.py

@@ -1,3 +1,6 @@
+import os
+os.environ["OMP_NUM_THREADS"] = "1"
+
 import gradio as gr
 import torch
 import re
@@ -56,6 +59,34 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 print("Model loaded successfully!")
 
+
+def format_output(xml_response):
+    def extract_field(field_name):
+        pattern = f'<{field_name}>(.*?)</{field_name}>'
+        match = re.search(pattern, xml_response, re.DOTALL)
+        return match.group(1).strip() if match else 'ERROR'
+    
+    argument = extract_field('argument')
+    claim = extract_field('claim')
+    explanation = extract_field('explanation')
+    verification = extract_field('human_verification_needed')
+    
+    if argument != 'NA' and argument != 'ERROR':
+        return f"""✅ **Argument Found**
+
+**Argument:** {argument}
+
+**Claim:** {claim}
+
+**Explanation:** {explanation}
+
+**Verification Needed:** {verification}"""
+    else:
+        return """❌ **No Argument Found**
+
+The text does not contain an argumentative unit."""
+
+
 @spaces.GPU
 def extract_arguments(text, temperature=0.1):
     if not text or not text.strip():
@@ -82,6 +113,7 @@ Extract arguments from historical text. Arguments in historical texts are often
             temperature = 0.3
     except:
         temperature = 0.1
+
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -96,8 +128,6 @@ Extract arguments from historical text. Arguments in historical texts are often
     generated_tokens = outputs[0][input_length:]
     response = tokenizer.decode(generated_tokens, skip_special_tokens=True)
 
-
-
     # Fix XML start
     if not response.startswith('<argument>'):
         arg_start = response.find('<argument>')
@@ -107,31 +137,6 @@ Extract arguments from historical text. Arguments in historical texts are often
     formatted = format_output(response)
     return response, formatted
 
-'''def format_output(xml_response):
-    def extract_field(field_name):
-        pattern = f'<{field_name}>(.*?)</{field_name}>'
-        match = re.search(pattern, xml_response, re.DOTALL)
-        return match.group(1).strip() if match else 'ERROR'
-    
-    argument = extract_field('argument')
-    claim = extract_field('claim')
-    explanation = extract_field('explanation')
-    verification = extract_field('human_verification_needed')
-    
-    if argument != 'NA' and argument != 'ERROR':
-        return f"""✅ **Argument Found**
-
-**Argument:** {argument}
-
-**Claim:** {claim}
-
-**Explanation:** {explanation}
-
-**Verification Needed:** {verification}"""
-    else:
-        return """❌ **No Argument Found**
-
-The text does not contain an argumentative unit."""'''
 
 # Gradio interface
 demo = gr.Interface(
@@ -153,6 +158,5 @@ demo = gr.Interface(
     ]
 )
 
-
 if __name__ == "__main__":
     demo.launch()