Update app.py

app.py

@@ -1,51 +1,46 @@
 import gradio as gr
-from unsloth import FastLanguageModel
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
 
-# 1. Load the model and adapters from the current folder
-# We use device_map="cpu" because the free Hugging Face tier doesn't have a GPU
-model, tokenizer = FastLanguageModel.from_pretrained(
-    model_name = ".", 
-    load_in_4bit = True,
-    device_map = "cpu"
+# 1. Load the Base Model and your Adapters using standard Transformers
+base_model_name = "unsloth/Llama-3.2-3B-bnb-4bit" # The base model
+adapter_path = "." # Your uploaded adapter files
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(adapter_path)
+
+# Load the base model in 4-bit (CPU compatible)
+model = AutoModelForCausalLM.from_pretrained(
+    base_model_name,
+    torch_dtype=torch.float32,
+    device_map="cpu",
+    low_cpu_mem_usage=True
 )
 
-# Enable fast inference mode
-FastLanguageModel.for_inference(model)
+# Merge your adapters onto the base model
+model = PeftModel.from_pretrained(model, adapter_path)
 
-# 2. Define the reasoning logic
+# 2. Reasoning Logic
 def legal_summarizer(legal_text):
-    # This template must match what you used during training
     prompt = f"Analyze the following legal text and provide a grounded summary.\n\nInput:\n{legal_text}\n\nResponse:\n"
+    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
     
-    # Tokenize and move to CPU
-    inputs = tokenizer([prompt], return_tensors = "pt").to("cpu")
-    
-    # Generate the response
-    outputs = model.generate(
-        **inputs, 
-        max_new_tokens = 256,
-        use_cache = True
-    )
-    
-    # Decode and clean the output to show ONLY the response
-    decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=256)
     
-    # Split by the 'Response:' tag to isolate the AI's answer
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
     if "Response:" in decoded:
         return decoded.split("Response:")[-1].strip()
     return decoded.strip()
 
-# 3. Setup the Gradio UI
+# 3. UI Setup
 demo = gr.Interface(
     fn=legal_summarizer, 
-    inputs=gr.Textbox(lines=10, label="Paste Legal Text/Contract", placeholder="Enter legal jargon here..."), 
+    inputs=gr.Textbox(lines=10, label="Paste Legal Text"), 
     outputs=gr.Textbox(label="LexGuard AI Summary"),
-    title="⚖️ LexGuard AI: Legal Reasoning Tool",
-    description="A fine-tuned Llama-3 model optimized for clarifying complex legal documents.",
-    theme="soft"
+    title="⚖️ LexGuard AI (CPU Edition)"
 )
 
-# Launch the app
 if __name__ == "__main__":
     demo.launch()