Update app.py

app.py

@@ -1,18 +1,15 @@
 import os
 os.environ["OMP_NUM_THREADS"] = "1"
 
-
 import gradio as gr
 import torch
 import re
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import spaces
 
-
 # Model configuration
 MODEL_ID = "oberbics/llama-3.1-newspaper-arguments-memi_2"
 
-
 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:
@@ -21,11 +18,12 @@ OUTPUT FORMAT - EXACTLY these 4 XML tags and NOTHING else:
 <explanation>Why this is an argument OR "NA"</explanation>
 <human_verification_needed>True OR False</human_verification_needed>
 
-<argument>It is reported from Malta that the British war-ships “Exmouth,” “Euryalus,” “Minerva,” and “Sutlej” have left for Messina. The French Government has sent two armoured ships and three destroyers to Messina. President Fallieres, Premier Clemenceau, Minister Pichon, and the Presidents of the Senate and Chamber have all sent messages of sympathy to the Italian Government. The help already proffered and accepted is insufficient for the purpose. There is pressing need of extraordinary measures of help, and provisions are in great demand. There is need of doctors, tents, clothing, and provisions for the survivors, who, deprived of all necessities, are exposed to the inclemencies of the winter weather. There is need of fire engines to cope with the flames that are raging among the ruins. The railway station has collapsed. Railway carriages have been destroyed. Almost all the railway employees are dead. The streets are no longer recognisable; they look like enormous fissures in a distant and extensive heap of ruins.</argument>
-<claim>Current relief efforts are inadequate and much more extensive aid is urgently needed.</claim>
-<explanation>The prefect explicitly argues that existing help is "insufficient" and makes a direct claim that "extraordinary measures" are needed, presenting a clear premise-conclusion structure about the inadequacy of current response.</explanation>
+<argument>Il giornale L'Italia moderna economica e finanziaria nel numero di oggi propone che non si facciano sottoscrizioni, le quali per quanto larghe sarebbero sempre impari ai bisogni, ma che il Parlamento stabilisca pochi centesimi addizionali per ogni lira su tutte le imposte e tasse (esclusi soltanto i dazi doganali la cui misura è vincolata da trattati di commercio).</argument>
+<claim>Private subscriptions are inadequate for earthquake relief; parliamentary taxation would be more effective.</claim>
+<explanation>The newspaper explicitly argues against private subscriptions as insufficient and proposes a specific alternative solution through parliamentary taxation, making a clear comparative argument about funding mechanisms.</explanation>
 <human_verification_needed>False</human_verification_needed>
 
+
 EXAMPLE WITHOUT ARGUMENT:
 <argument>NA</argument>
 <claim>NA</claim>
@@ -52,77 +50,102 @@ print("Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 tokenizer.pad_token = tokenizer.eos_token
 
-
 print("Loading 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
-)
-
-
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    quantization_config=bnb_config,
-    device_map="auto",
-    torch_dtype=torch.bfloat16,
-    trust_remote_code=True
-)
-print("Model loaded successfully!")
-
-
-
-
-
-
 
+# Try to load with bfloat16 first (best quality)
+# If that fails due to memory, fall back to 8-bit quantization
+try:
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+        trust_remote_code=True
+    )
+    print("✅ Model loaded successfully in bfloat16!")
+except Exception as e:
+    print(f"⚠️ Could not load in bfloat16: {e}")
+    print("Trying 8-bit quantization...")
+    
+    try:
+        # 8-bit is much better than 4-bit
+        bnb_config = BitsAndBytesConfig(
+            load_in_8bit=True,
+            llm_int8_threshold=6.0
+        )
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            quantization_config=bnb_config,
+            device_map="auto",
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True
+        )
+        print("✅ Model loaded successfully in 8-bit!")
+    except Exception as e:
+        print(f"⚠️ Could not load in 8-bit: {e}")
+        print("Falling back to 4-bit quantization (may reduce quality)...")
+        
+        # Last resort: 4-bit quantization
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16,
+            bnb_4bit_use_double_quant=True
+        )
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            quantization_config=bnb_config,
+            device_map="auto",
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True
+        )
+        print("⚠️ Model loaded in 4-bit (quality may be reduced)")
 
 @spaces.GPU
 def extract_arguments(text, temperature=0.1):
     if not text or not text.strip():
-        return "", "Please enter some text to analyze."
+        return "Please enter some text to analyze."
+    
+    # Use the same message format as the working local version
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": f"Extract argumentative units from historical text in their original form, no summary.\n{text}"}
+    ]
     
-    prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
-{SYSTEM_PROMPT}<|eot_id|>
-<|start_header_id|>user<|end_header_id|>
-Extract argumentative units from historical text in their original form, no summaries. Factual reportings such as "Dem Vulkanausbruch folgten drei Sturzwellen in etwa 10 Meter Höhe" are no arguments
-{text}<|eot_id|>
-<|start_header_id|>assistant<|end_header_id|>"""
+    # Use apply_chat_template - the proper way to format prompts
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
     
-    inputs = tokenizer(prompt, return_tensors="pt",truncation=True, max_length=5048).to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=5048).to(model.device)
     input_length = inputs['input_ids'].shape[1]
+    
+    # Validate temperature
     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
+        temperature = float(temperature)
+        temperature = max(0.01, min(0.3, temperature))
     except:
         temperature = 0.1
-
-
+    
+    # Generate with the same parameters as local version
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
             max_new_tokens=8000,
             temperature=temperature,
-            do_sample=True if temperature > 0.01 else False,
-            top_p=0.9,
-            pad_token_id=tokenizer.eos_token_id,
-            repetition_penalty=1.1
+            top_p=0.95,  # Match local version
+            repetition_penalty=1.15,  # Match local version
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
         )
     
     generated_tokens = outputs[0][input_length:]
     response = tokenizer.decode(generated_tokens, skip_special_tokens=True)
-
-
-    # Fix XML start
+    
+    # Fix XML start if needed
     if not response.startswith('<argument>'):
         arg_start = response.find('<argument>')
         if arg_start != -1:
@@ -130,9 +153,6 @@ Extract argumentative units from historical text in their original form, no summ
     
     return response
 
-
-
-
 # Gradio interface
 demo = gr.Interface(
     fn=extract_arguments,
@@ -147,13 +167,9 @@ demo = gr.Interface(
     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]
+        ["England ist nicht bereit lür die Olympiade 1948 » sagt Vernon Morgan, der frühere Olympialäufer und südamerikanische Rekordler über 1500 m. Wir geben nachstehendem Artikel noch Raum, ob- derer Stelle von größerem Nutzen; alle Baumatewohl es inzwischen keinem Zweifel mehr unter- rialien würden dringend zur Behebung der kataliegen wird, daß England die Olympiade durch- strophalen Wohnungsnot gebraucht, und die Nahführen wird, wenn es auch noch große Schwierig- rungsmittelversorgung sei ohnehin schon knapp gekeiten überwinden muß. nug und er) aube eine soicne Rjeseneinladung nicht. Im kommenden Sommer werden in London Sportler, Dann gibt es solche, die überhaupt gegen die offizielle Würdenträger und Besucher aus aller Welt Olympischen Spiele sind, mit der Begründung, dazu dem zweifellos größten sportlichen Ereignis der Positive, was an ihnen sei, werde bei weitem durch Welt zusammenkommen— zu den Olympischen Spie- das Negative übertroffen, und sie führten nur zu len. Seit vierzig Jahren, nämlich seit die Olymplade Internationalen Zänkereien. von 1908 in der alten White City abgehalten wurde, Dies sind in großen Zügen die verschiedenen Anhat London nicht mehr eine solche Ansammlung von sichten. Und ich als glühender Verehrer der olym-Sportsleuten aller Rassen, Farben und Glaubensbe- pischen Bewegung und als jemand, der im Sport und kenntnisse in Englands Hauptstadt gesehen. in der Sportlichkeit eines der Mittel sieht, der Welt Das erste, was heute jeder sagt, wenn er von der den Frieden zu bringen — auch Ich bin der Meinung, Londoner Olympiade spricht, ist: « Soll England sie daß die Olympischen Spiele im nächsten Jahr nicht abhalten?» Wahrscheinlich wird diese Frage ebenso in London abgehalten werden sollten, häufig in den USA. und Indien — um nur zwei große Daß alles nur Mögliche, was für unsere J _ OSUiher Länder zu nennen — wie In England gestellt. Meldungen aus Uebersee besagen, daß im Ausland über diese Frage ebenso viel Meinungsverschiedenheit herrscht wie in England. Im Vereinigten Königreich werden folgende Argumente geltend gemacht: « Es ist eine große Ehre für ein Land und seine Hauptstadt, wenn ihm die Olympischen Spiele zugewiesen werden. Die Nationen der Welt haben die Olympiade von 1948 London übergeben; es abzuweisen, wäre unhöflich — nicht davon zu reden, daß viele Jahre darüber hingehen könnten, bis England wieder eine solche Ehre zuteil werden würde. » Viel zu wenig Menschen in der Welt sind sich aber darüber klar, wie sich das Leben heute in England abspielt, und die Besucher der Olymplade werden selbst imstande sein zu sehen, wie hart die m Krieg heimgesuchten Einwohner Großbritan < iii um den Frieden kämpfen. England ist nicht bereit für die Besucher Diejenigen, die gegen eine Abhaltung der Spiele in London sind, sagen als erstes, daß das Land nicht darauf vorbereitet sei, eine so umfangreiche sportliche Festivität abzuhalten, und die Besucher würden enttäuscht und unzufrieden mit ihren Gastgebern nach Hause zurückkehren. « Lieber niemand ins Haus einladen », sagen sie, « als ihm ein miserables Essen vorsetzen ». Sie sagen ferner, das ganze Geld, das für die Spiele ausgegeben werden soll, auch wenn die Vorführungen in einem bescheidenen uneinfachen Rahmen abgehalten werden, wäre _ an an getan werden kann, auch getan werden wird, darüber isl kein Wort zu verlieren; aber unser Land ist nicht imstande, einen Zufluß von Tausenden von Besuchern zu irgendwelchem Zweck aufzunehmen. Es ist viel zu wenig von allem da. Die Nationen waren durch die Annahme überrascht Wenn London rechtzeitig gebeten hätte, von ihm abzusehen, würde es bei den übrigen teilnehmenden Nationen volles Verständnis gefunden haben. Es gab maßgebende Leute, die über die Annahme Englands ziemlich überrascht waren. Von einem Gastgeber, der sein Diner wegen mangelnder Vorräte rechtzeitig absagt, denkt man oft besser als von einem, der keine Vorsorge trifft, und seine Gäste bittet, mit lern vorlieb zu nehmen, was gerade da ist. Es ist menschlich, nicht mit allen Begleitumständen einverstanden zu sein, und ein solcher Gastgeber könnte wohl der Ungasllichkeit geziehen werden. Anderseits kann man allerdings nicht umhin zuzugeben, daß es Ausländern gut tüte, zu sehen, wie England heute lebt. Was das Argument betrifft, daß bei den Spielen mehr Negatives als Gutes herauskäme, so ist hier nicht der rechte Ort, diese vielumstrittene und dornige Frage zu untersuchen, aber man kann zum mindesten sagen, daß die Mehrheit aller Leute in der Welt ihre Abhaltung befürwortet. Etwa vor einem Jahr stellte der « Daily Express », die am meisten verbreitete englische Tageszeitung, eine Rundfrage darüber an, ob England die Spiele Im nach", 0.1]
     ]
 )
 
-
-
-
 if __name__ == "__main__":
-    demo.launch() 
-
+    demo.launch()