Add application file

app.py

@@ -2,27 +2,51 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from peft import PeftModel
 
-# Load base + LoRA fine-tuned model
-base_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-adapter_model = "Noshitha98/TinyLlama-ToS-Finetuned"
+# Global cache for the pipeline
+pipe = None
 
-tokenizer = AutoTokenizer.from_pretrained(base_model)
-model = AutoModelForCausalLM.from_pretrained(base_model)
-model = PeftModel.from_pretrained(model, adapter_model)
+def load_model():
+    """Load the TinyLlama model with LoRA adapters (cached)."""
+    global pipe
+    if pipe is None:
+        base_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+        adapter_model = "Noshitha98/TinyLlama-ToS-Finetuned"
 
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+        tokenizer = AutoTokenizer.from_pretrained(base_model)
+        model = AutoModelForCausalLM.from_pretrained(base_model)
+        model = PeftModel.from_pretrained(model, adapter_model)
+
+        pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+    return pipe
 
 def check_clause(text):
-    # Instruction prompt for anomaly detection
+    generator = load_model()
+
+    if not text.strip():
+        return [("⚠️ Please enter contract text.", None)]
+
     prompt = f"Detect unusual or risky clauses in this contract:\n\n{text}\n\nResponse:"
-    output = pipe(prompt, max_length=512, do_sample=False)[0]["generated_text"]
-    return output
+    output = generator(prompt, max_length=512, do_sample=False)[0]["generated_text"]
+
+    # Clean the response
+    response = output.split("Response:")[-1].strip()
+
+    # Simple heuristic: highlight sentences with "risk", "penalty", "terminate", "discretion"
+    risky_keywords = ["penalty", "terminate", "discretion", "risk", "sole", "modify", "liability"]
+    highlights = []
+    for sentence in text.split(". "):
+        label = "risky" if any(word.lower() in sentence.lower() for word in risky_keywords) else None
+        highlights.append((sentence.strip(), label))
+
+    return highlights
 
-# Gradio UI
+# Gradio UI with highlighted output
 iface = gr.Interface(
     fn=check_clause,
     inputs=gr.Textbox(lines=8, placeholder="Paste contract clause here..."),
-    outputs="text",
+    outputs=gr.HighlightedText(
+        color_map={"risky": "red"}
+    ),
     title="Contract Clause Checker",
     description="Paste contract text to detect unusual or anomalous clauses using TinyLlama-ToS-Finetuned."
 )