Update app.py

app.py

@@ -1,47 +1,78 @@
-# -*- coding: utf-8 -*-
-"""app.ipynb
-
-Automatically generated by Colab.
-
-Original file is located at
-    https://colab.research.google.com/drive/1ewSnt_PQx8cX9GC2XAguVNesbn8MM3tA
+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".
 """
 
-
-
-from fastapi import FastAPI
-from pydantic import BaseModel
-from transformers import RobertaTokenizer, RobertaForMaskedLM, pipeline
-
-app = FastAPI()
-
-
-model_name = "web3se/SmartBERT-v3"
-tokenizer = RobertaTokenizer.from_pretrained(model_name)
-model = RobertaForMaskedLM.from_pretrained(model_name)
-
-# Create fill-mask pipeline
-fill_mask = pipeline('fill-mask', model=model, tokenizer=tokenizer)
-
-# Define request body
-class ContractRequest(BaseModel):
-    contract_code: str
-
-@app.post("/analyze/")
-async def analyze_contract(request: ContractRequest):
-    # Use fill-mask on contract code
-    outputs = fill_mask(request.contract_code)
-
-    return {"predictions": outputs}  # Returns top masked token predictions
-
-@app.get("/")
-async def root():
-    return {"message": "FastAPI is running!"}
-
-
-# Run API on Hugging Face Spaces
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
-
-
+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()