Update app.py

app.py

@@ -1,19 +1,16 @@
-# app.py
 import os
 import json
 import warnings
-import re
 import torch
 import gradio as gr
 import spaces
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from transformers import StoppingCriteria, StoppingCriteriaList
 
 # ---------- CONFIG ----------
 os.environ.setdefault("GRADIO_SERVER_PORT", "7860")
-MODEL_PATH = "iqasimz/g1"  # <- change to your repo or local dir
+MODEL_PATH = "iqasimz/g2"  # <- change to your repo or local dir
 MAX_NEW_TOKENS_DEFAULT = 300
-TEMPERATURE_DEFAULT = 0.2
+TEMPERATURE_DEFAULT = 0
 TOP_P_DEFAULT = 1.0
 # ---------------------------
 
@@ -46,12 +43,12 @@ def load_model_to_cpu(model_dir: str):
 
 def build_inference_prompt(paragraph: str) -> str:
     # Match your training format EXACTLY (Task + Rules + Paragraph in user turn)
-    task_block = """Task: You are an expert argument analyst. Number the sentences in the paragraph and tag the role of each one.
-Rules:
-- Do NOT change the text of any sentence.
-- Keep the original order.
-- Output exactly N lines, one per sentence.
-- Each line must be: "<index> <original sentence> <role>", where role ∈ {claim, premise, none}.
+    task_block = """Task: ou are an expert argument analyst. Identify the role of each sentence within the context of the paragraph/debate/article like a true linguistics and argument expert.Number the sentences in the paragraph and tag the role of each one.\n
+Rules:\n
+- Do NOT change the text of any sentence.\n
+- Keep the original order.\n
+- Output exactly N lines, one per sentence.\n
+- Each line must be: "<index> <original sentence> <role>", where role ∈ {claim, premise, none}.\n
 - Do not add any explanations or extra text after the Nth line.
 """
     # Chat-style formatting used during training
@@ -60,71 +57,111 @@ Rules:
         f"<|im_end|>\n<|im_start|>assistant\n"
     )
 
-# -------- Sentence counting for N --------
-SENT_SPLIT_RE = re.compile(r'(?<!\b[A-Z])(?<=[.!?])\s+(?=\S)')
+def get_last_five_words(text: str) -> str:
+    """Get the last 5 words from a text string."""
+    words = text.strip().split()
+    return " ".join(words[-5:]) if len(words) >= 5 else " ".join(words)
 
-def count_sentences(paragraph: str) -> int:
-    p = (paragraph or "").strip()
-    if not p:
-        return 0
-    parts = [s.strip() for s in SENT_SPLIT_RE.split(p) if s.strip()]
-    return max(1, len(parts))
-
-# -------- Stopping criteria to halt after N labeled lines --------
-class RoleLinesStop(StoppingCriteria):
+def extract_role_from_suffix(text_after_match: str) -> str:
     """
-    Stop when we've generated N lines that look like:
-      <index> <original sentence> <role>
-    with role ∈ {claim, premise, none}.
-    Also stops if the model begins line N+1 (e.g., "N+1 ").
+    Extract role (claim, premise, none) from text after the 5-word match.
+    Handles cases like 'claimabcd' -> 'claim'
     """
-    def __init__(self, tokenizer, prompt_len: int, n_lines: int):
-        self.tok = tokenizer
-        self.prompt_len = prompt_len
-        self.n_lines = n_lines
-        self.role_line_re = re.compile(
-            r'^\s*\d+\s+.+\s+(?:claim|premise|none)\s*$', re.IGNORECASE | re.MULTILINE
-        )
-        self.next_index_re = re.compile(rf'^\s*{n_lines+1}\s', re.MULTILINE) if n_lines >= 1 else None
-
-    def __call__(self, input_ids, scores, **kwargs) -> bool:
-        gen_ids = input_ids[0, self.prompt_len:]
-        if gen_ids.numel() == 0:
-            return False
-        text = self.tok.decode(gen_ids, skip_special_tokens=True)
-
-        # If we see the start of line N+1, stop immediately
-        if self.next_index_re and self.next_index_re.search(text):
-            return True
-
-        # Count complete role-tagged lines
-        complete_lines = self.role_line_re.findall(text)
-        return len(complete_lines) >= self.n_lines
-
-def parse_numbered_lines(text: str):
+    text_after_match = text_after_match.strip()
+    
+    # Look for the role words at the start of the remaining text
+    role_words = ['claim', 'premise', 'none']
+    for role in role_words:
+        if text_after_match.lower().startswith(role.lower()):
+            return role
+    
+    # If no exact match, return the first word (fallback)
+    first_word = text_after_match.split()[0] if text_after_match.split() else ""
+    for role in role_words:
+        if first_word.lower().startswith(role.lower()):
+            return role
+    
+    return "none"  # default fallback
+
+def parse_numbered_lines(text: str, original_paragraph: str):
     """
-    Optional: parse lines like:
-    1 Some sentence. claim
-    2 Another sentence. premise
-    into a list of dicts.
+    Enhanced parsing with improved stopping criteria:
+    1. Find exact match of last 5 words from input paragraph
+    2. Look for role word after a space following the match
+    3. Stop parsing after finding the last sentence to avoid gibberish
     """
     results = []
-    for line in text.splitlines():
+    lines = text.splitlines()
+    
+    # Get sentences from original paragraph for reference
+    import re
+    sentences = re.split(r'[.!?]+', original_paragraph.strip())
+    sentences = [s.strip() for s in sentences if s.strip()]
+    
+    if not sentences:
+        return results
+    
+    # Get last 5 words of the original paragraph
+    last_five_words = get_last_five_words(original_paragraph)
+    
+    for line in lines:
         line = line.strip()
         if not line or not line[0].isdigit():
             continue
+            
         try:
-            # index first
+            # Parse index
             space_after_idx = line.find(" ")
+            if space_after_idx == -1:
+                continue
+                
             idx = int(line[:space_after_idx])
             rest = line[space_after_idx + 1:].rstrip()
-            # last space => role
-            last_space = rest.rfind(" ")
-            sent = rest[:last_space].strip()
-            role = rest[last_space + 1:].strip().lower()
-            results.append({"index": idx, "sentence": sent, "role": role})
-        except Exception:
-            pass
+            
+            # Check if this line contains the last 5 words (indicating last sentence)
+            if last_five_words.lower() in rest.lower():
+                # Find the position of the last 5 words
+                match_pos = rest.lower().find(last_five_words.lower())
+                if match_pos != -1:
+                    # Extract sentence (everything up to and including the match)
+                    sentence_end = match_pos + len(last_five_words)
+                    sent = rest[:sentence_end].strip()
+                    
+                    # Look for role after the match
+                    text_after_match = rest[sentence_end:].strip()
+                    role = "none"  # default
+                    
+                    if text_after_match:
+                        # Skip any immediate punctuation/spaces and look for role
+                        text_after_match = text_after_match.lstrip(' .,!?')
+                        role = extract_role_from_suffix(text_after_match)
+                    
+                    results.append({"index": idx, "sentence": sent, "role": role})
+                    
+                    # STOP parsing here - this is the last sentence
+                    break
+            else:
+                # Regular parsing for non-last sentences
+                last_space = rest.rfind(" ")
+                if last_space == -1:
+                    continue
+                    
+                sent = rest[:last_space].strip()
+                role_candidate = rest[last_space + 1:].strip().lower()
+                
+                # Clean role (handle gibberish suffixes)
+                role = "none"
+                for valid_role in ['claim', 'premise', 'none']:
+                    if role_candidate.startswith(valid_role):
+                        role = valid_role
+                        break
+                
+                results.append({"index": idx, "sentence": sent, "role": role})
+                
+        except Exception as e:
+            print(f"Error parsing line '{line}': {e}")
+            continue
+    
     return results
 
 @spaces.GPU(duration=120)
@@ -139,26 +176,16 @@ def analyze(paragraph: str, max_new_tokens: int, temperature: float, top_p: floa
     prompt = build_inference_prompt(paragraph)
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
 
-    # Compute target number of lines (N) and install stopping criteria
-    n_lines = count_sentences(paragraph)
-    stopper = RoleLinesStop(
-        tokenizer=tokenizer,
-        prompt_len=inputs["input_ids"].shape[1],
-        n_lines=n_lines
-    )
-    stops = StoppingCriteriaList([stopper])
-
     with torch.inference_mode():
         output = model.generate(
             **inputs,
             max_new_tokens=int(max_new_tokens),
             temperature=float(temperature),
             top_p=float(top_p),
-            do_sample=(float(temperature) > 0.0),  # sampling only if temp > 0
+            do_sample=(temperature > 0.0 and top_p < 1.0),
             pad_token_id=tokenizer.pad_token_id,
             eos_token_id=tokenizer.eos_token_id,
             use_cache=True,
-            stopping_criteria=stops,
         )
 
     full = tokenizer.decode(output[0], skip_special_tokens=False)
@@ -170,15 +197,8 @@ def analyze(paragraph: str, max_new_tokens: int, temperature: float, top_p: floa
     else:
         resp = full.strip()
 
-    # Safety net: hard-trim to exactly N labeled lines if model leaked extras
-    role_line_re = re.compile(r'^\s*\d+\s+.+\s+(?:claim|premise|none)\s*$', re.IGNORECASE | re.MULTILINE)
-    matched = role_line_re.findall(resp)
-    if matched:
-        trimmed = "\n".join(matched[:n_lines]).strip()
-        if trimmed:
-            resp = trimmed
-
-    parsed = parse_numbered_lines(resp)
+    # Updated parsing with original paragraph reference
+    parsed = parse_numbered_lines(resp, paragraph)
     parsed_json = json.dumps(parsed, ensure_ascii=False, indent=2) if show_parsed else ""
     return resp, parsed_json
 
@@ -218,9 +238,10 @@ def launch_app():
         )
 
         gr.Markdown("### Tips")
-        gr.Markdown("- Set `MODEL_PATH` at the top to your merged model repo or local path.\n"
+        gr.Markdown("- Set MODEL_PATH at the top to your merged model repo or local path.\n"
                     "- For deterministic outputs, set Temperature=0.0 and Top-p=1.0.\n"
-                    "- Output is forcibly stopped after exactly N lines.")
+                    "- Your training format (chat tokens + Task/Rules) is preserved in the prompt.\n"
+                    "- **Enhanced parsing**: Stops at last sentence using 5-word match to avoid gibberish.")
 
     return demo