File size: 2,945 Bytes
cd1f97c 6a70183 cd1f97c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
# Load model and tokenizer
@gr.on_startup()
def load_model():
global tokenizer, model
tokenizer = AutoTokenizer.from_pretrained("web3se/SmartBERT-v3")
model = AutoModelForSequenceClassification.from_pretrained("web3se/SmartBERT-v3")
model.eval()
return model, tokenizer
# Function to analyze smart contract
def analyze_contract(contract_code):
if not contract_code:
return {"error": "No contract code provided"}
try:
# Tokenize input
inputs = tokenizer(contract_code, return_tensors="pt", truncation=True, max_length=512)
# Perform inference
with torch.no_grad():
outputs = model(**inputs)
# Process outputs
logits = outputs.logits
probabilities = torch.nn.functional.softmax(logits, dim=-1)
# Get predicted class and confidence
predicted_class = torch.argmax(probabilities, dim=-1).item()
confidence = probabilities[0, predicted_class].item()
# Map class index to label
labels = model.config.id2label
predicted_label = labels[predicted_class]
# Format all class probabilities for display
all_probs = {labels[i]: f"{prob.item()*100:.2f}%" for i, prob in enumerate(probabilities[0])}
sorted_probs = dict(sorted(all_probs.items(), key=lambda item: float(item[1].rstrip('%')), reverse=True))
# Format the result as markdown for better display
result_md = f"## Analysis Results\n\n"
result_md += f"**Prediction:** {predicted_label}\n\n"
result_md += f"**Confidence:** {confidence*100:.2f}%\n\n"
result_md += "### All Class Probabilities:\n\n"
for label, prob in sorted_probs.items():
result_md += f"- {label}: {prob}\n"
return result_md
except Exception as e:
return f"Error: {str(e)}"
# Create Gradio interface
title = "Smart Contract Analyzer (SmartBERT-v3)"
description = """
This app uses web3se/SmartBERT-v3 model to analyze smart contracts.
Simply paste your smart contract code in the text area below and click "Analyze".
"""
iface = gr.Interface(
fn=analyze_contract,
inputs=gr.Textbox(lines=15, placeholder="Paste your smart contract code here...", label="Smart Contract Code"),
outputs=gr.Markdown(label="Analysis Results"),
title=title,
description=description,
examples=[
# You can add example smart contracts here
["pragma solidity ^0.8.0;\n\ncontract SimpleStorage {\n uint256 private value;\n \n function set(uint256 _value) public {\n value = _value;\n }\n \n function get() public view returns (uint256) {\n return value;\n }\n}"]
],
allow_flagging="never"
)
# Launch the app
iface.launch() |