Update app.py

app.py

@@ -3,137 +3,94 @@ os.environ["OMP_NUM_THREADS"] = "1"
 
 import gradio as gr
 import torch
-import re
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from peft import PeftModel
 import spaces
 
-# Model configuration
-MODEL_ID = "oberbics/llama-3.1-newspaper-arguments-your_name-full"
-
-SYSTEM_PROMPT = """You are an expert at analyzing historical texts and you hate to summarize
-
-OUTPUT FORMAT - EXACTLY these 4 XML tags and NOTHING else:
-<argument>Original argument text OR "NA"</argument>
-<claim>Core claim (implication) in one sentence OR "NA"</claim>
-<explanation>Why this is an argument OR "NA"</explanation>
-<human_verification_needed>True OR False</human_verification_needed>
-
-EXAMPLE WITH ARGUMENT:
-<argument>Es sind furchtbare Bilder, die sich dabei entrollen. Unter den Trümmern des einen Hause», so erzählt Luigt Barsint im Corrtcre della sera, findet man die Leichen von Unglück lichen, die in anderen Häusern gewohnt baben und die in der Ber- Wirrung de» schrcck.ichen Augenblickes instinktiv bet Fremden Hülfe und Unterschlupf suchten. Niemand erkennt jetzt diese armen Ein dringlinge, ihre Leichen werden nicht reklamiert, und man trägt sie hinunter an de» Strand, wo sie in langer Reihe einer neben den anderen hingebettet weiden, in denselben Tüchern und Decken, in denen sie tbren Tod gesunden.</argument>
-<claim>The earthquake's chaos led to unidentified victims dying in unfamiliar places.</claim>
-<explanation>Describes how people fled to other houses seeking help during the disaster, died there, and now cannot be identified or claimed by relatives. Shows cause (panic/confusion) and effect (anonymous deaths).</explanation>
-<human_verification_needed>False</human_verification_needed>
-
-EXAMPLE WITHOUT ARGUMENT:
-<argument>NA</argument>
-<claim>NA</claim>
-<explanation>NA</explanation>
-<human_verification_needed>FALSE</human_verification_needed>
-
-RULES:
-- NO SUMMARY; ONLY ORIGINAL EXTRACTOM FROM THE TEXT; don't extract anything that is not in the text. Only extract word by word
-- ONLY output these 4 XML tags
-- Factual reportings such as "Dem Vulkanausbruch folgten drei Sturzwellen in etwa 10 Meter Höhe" or "Almost all the inhabitants were killed; only a few thousands escaped death" are NO arguments
-- Extract only original text without changes or use NA when you did not find an argument
-- The claim is not a translation of summary of argument. It should say what the (implicite) argument implies
-- In cases of uncertainty or ambiguity, say human_verification_needed TRUE
-- If no argument exists, use NA for all fields except <human_verification_needed>FALSE or TRUE</human_verification_needed>
-- More than one argumentative unit possible for one aticle, one unit has one clear clame and all the xml structures
-
-VERIFICATION: BEfore you print the results, double check explanations of the argument. When the explanation states that this is not an argument, don't print it"""
+# Base and adapter
+BASE_MODEL = "meta-llama/Meta-Llama-3.3-70B-Instruct"   # official base model
+ADAPTER = "oberbics/llama-3.3-70B-adapter-newspaper-argument-mining"
+
+SYSTEM_PROMPT = """... your system prompt ..."""
 
 print("Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 tokenizer.pad_token = tokenizer.eos_token
 
-print("Loading model...")
+print("Loading base model...")
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
     bnb_4bit_compute_dtype=torch.bfloat16,
-    bnb_4bit_use_double_quant=True
+    bnb_4bit_use_double_quant=True,
 )
 
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
     quantization_config=bnb_config,
     device_map="auto",
     torch_dtype=torch.bfloat16,
     trust_remote_code=True
 )
-print("Model loaded successfully!")
 
+print("Loading adapter...")
+model = PeftModel.from_pretrained(base_model, ADAPTER)
 
+# (Optional) merge LoRA weights into the base model for faster inference
+# model = model.merge_and_unload()
 
+print("Model + adapter loaded successfully!")
 
 @spaces.GPU
 def extract_arguments(text, temperature=0.1):
     if not text or not text.strip():
         return "", "Please enter some text to analyze."
-    
+
     prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
 {SYSTEM_PROMPT}<|eot_id|>
 <|start_header_id|>user<|end_header_id|>
 Extract arguments from historical text.
 {text}<|eot_id|>
 <|start_header_id|>assistant<|end_header_id|>"""
-    
+
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=10048).to(model.device)
     input_length = inputs['input_ids'].shape[1]
-    try:
-        if temperature is None:
-            temperature = 0.1
-        else:
-            temperature = float(temperature)
-
-        if temperature < 0.01:
-            temperature = 0.01
-        elif temperature > 0.3:
-            temperature = 0.3
-    except:
-        temperature = 0.1
+
+    temperature = float(temperature) if temperature else 0.1
+    temperature = max(0.01, min(temperature, 0.3))
 
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
             max_new_tokens=5000,
             temperature=temperature,
-            do_sample=True if temperature > 0.01 else False,
+            do_sample=temperature > 0.01,
             top_p=0.9,
             pad_token_id=tokenizer.eos_token_id,
-            repetition_penalty=1.1
+            repetition_penalty=1.1,
         )
-    
+
     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>')
         if arg_start != -1:
             response = response[arg_start:]
-    
-    return response
 
+    return response
 
-# Gradio interface
+# Gradio UI
 demo = gr.Interface(
     fn=extract_arguments,
     inputs=[
         gr.Textbox(label="Input Text", placeholder="Enter newspaper text here...", lines=10),
         gr.Slider(minimum=0.01, maximum=0.3, value=0.1, step=0.01, label="Temperature (lower = more consistent)")
     ],
-    outputs=[
-        gr.Textbox(label="Raw XML Output", lines=8)
-    ],
+    outputs=[gr.Textbox(label="Raw XML Output", lines=8)],
     title="Newspaper Argumentative Unit Extractor",
-    description="Extract argumentative units from news sources",
-    examples=[
-        ["Reggio, January 8. Frequent shocks of earthquake were felt here dur ing the night, accompanied at times by loud subter ranean reports. A few buildings that had not been completely destroyed were further damaged. The work of reconstructing the railway is being pushed forward energetically. News has been received from Brancaleone, Catanzaro, and Palmi of earthquakes by which the inhabitants were alarmed last night", 0.1],
-        ["Rome, January 20. At the request of the Queen, Dr. Quinco, the Court physician, has left Rome in order to super vise the medical arrangements, and the distribution of clothing and linen which has been prepared at the Quirinal under Her Majesty’s eye, in places on the Calabrian coast and inland where medical help is scarce. Rome, January 20. The controversy as to whether Messina is to be rebuilt on its former lines or finally abandoned is still in full swing. The newspapers state that a mere village will probably represent Messina for many years to come, as even should the rebuilding project eventually be decided upon, the work of clearing away the enormous quantity of debris and constructing new foundations must extend over a lengthy period. Plans for the erection of so-called earthquake-proof houses are engaging the attention of architects, but pressmen who visited the scene shortly after the disaster state that one building of this description was utterly wrecked. The Mu- sella plain, not far distant from Messina, has been selected as a site for the erection of temporary habi tations, after having been inspected by Signor Ca- sana, the Minister for War. The problem of pro viding temporary accommodation becomes daily more pressing, as the majority of the survivors refuse to leave their household goods lying under the ruins to the mercy of plunderers. Both Catania and Pa lermo, however, are crowded with refugees, and great difficulty is experienced in finding shelter for them. General Mazza, the provisional commandant at Messina, continues to send favourable reports regarding the health of the troops and survivors. Conditions at Reggio remain the same, but no fur ther rescues have been reported from there since Saturday. Milan, January 19. Signor Barzini, writing £n the Corriere de la Sera, accuses the authorities at Messina of inactivity, want of energy, and pedantry.", 0.1]
-    ]
+    description="Extract argumentative units from news sources"
 )
 
-
 if __name__ == "__main__":
     demo.launch()