Done

app.py

@@ -3,16 +3,17 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from fpdf import FPDF
 from datetime import date
 import os
-import re
 
-# Load Hugging Face model
+# Load model with Hugging Face token
 token = os.getenv("HF_TOKEN")
-model_id = "syedMohib44/ai-auditor-model-l"
+model_id = "syedMohib44/ai-auditor-model"
 
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
 model = AutoModelForCausalLM.from_pretrained(model_id, token=token)
+
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
+# PDF generator
 def generate_audit_pdf(code, description, findings):
     pdf = FPDF()
     pdf.add_page()
@@ -24,7 +25,7 @@ def generate_audit_pdf(code, description, findings):
 
     pdf.set_font("Arial", '', 12)
     pdf.cell(0, 10, f"Scan Date: {date.today().strftime('%Y-%m-%d')}", ln=True)
-    pdf.cell(0, 10, "Model: Gemma 2B IT with LoRA", ln=True)
+    pdf.cell(0, 10, "Model: TinyLlama-1.1B-Chat with LoRA", ln=True)
     pdf.cell(0, 10, "Audit Engine: AI Audit Agent", ln=True)
     pdf.ln(5)
 
@@ -55,56 +56,25 @@ def generate_audit_pdf(code, description, findings):
     pdf.output(output_path)
     return output_path
 
-def extract_sections(text):
-    headers = ["High Severity", "Medium Severity", "Low Severity", "Best Practices"]
-    findings = {}
-
-    for i, header in enumerate(headers):
-        if i + 1 < len(headers):
-            next_header = headers[i + 1]
-            pattern = rf"{header}:(.*?)(?=\n{next_header}:)"
-        else:
-            pattern = rf"{header}:(.*?)(?=\nDisclaimer:|\Z)"
-
-        match = re.search(pattern, text, re.DOTALL | re.IGNORECASE)
-        findings[header] = match.group(1).strip() if match else "No issues found."
-
-    return findings
-
-
+# Inference function for Gradio
 def audit_contract(code, description):
-    prompt = f"""
-### Instruction:
-You are a professional smart contract auditor.
-
-Audit the following smart contract and generate a structured report.
-
-Each severity level should contain:
-- Title
-- Description
-- Recommendation
-
-If no issues are found in a category, return: "No issues found."
-
-
+    prompt = f"""### Instruction:
+Please audit the following smart contract and provide a recommendation.
+### Code:
+{code}
 ### Description:
 {description}
-
-### Audit Report:
-
-Disclaimer: This is an automated audit report generated by AI. Manual review by security experts is recommended.
+### Recommendation:
 """
-    response = pipe(prompt, max_new_tokens=800, do_sample=False)[0]["generated_text"]
-    
-    # Print raw response for debugging
-    print("=== RAW MODEL OUTPUT ===")
-    print(response)
+    result = pipe(prompt, max_new_tokens=300)[0]["generated_text"]
+    findings = result.split("### Recommendation:")[-1].strip()
+
+    # Generate PDF
+    pdf_path = generate_audit_pdf(code, description, findings)
 
-    findings_dict = response #extract_sections(response)
-    pdf_path = generate_audit_pdf(code, description, findings_dict)
-    # combined_findings = "\n\n".join(f"{k}:\n{v}" for k, v in findings_dict.items())
-    return response, pdf_path
+    return findings, pdf_path
 
+# Gradio UI
 iface = gr.Interface(
     fn=audit_contract,
     inputs=[
@@ -116,8 +86,8 @@ iface = gr.Interface(
         gr.File(label="Download Audit Report")
     ],
     title="AI Smart Contract Auditor",
-    description="Paste your Solidity contract and description. The AI will generate a structured audit with severity levels and a downloadable PDF."
+    description="Paste your smart contract code and description. The AI will generate an audit with findings and a downloadable PDF report."
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()