Update app.py

app.py

@@ -1,78 +1,25 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import RobertaTokenizer, RobertaForMaskedLM, pipeline
 
-# 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
+app = FastAPI()
 
-# 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)}"
+# Load SmartBERT v3
+model_name = "web3se/SmartBERT-v3"
+tokenizer = RobertaTokenizer.from_pretrained(model_name)
+model = RobertaForMaskedLM.from_pretrained(model_name)
 
-# 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".
-"""
+# Define API input format
+class ContractRequest(BaseModel):
+    contract_code: str
 
-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"
-)
+@app.post("/analyze/")
+async def analyze_contract(request: ContractRequest):
+    fill_mask = pipeline('fill-mask', model=model, tokenizer=tokenizer)
+    outputs = fill_mask(request.contract_code)
+    return {"predictions": outputs} 
 
-# Launch the app
-iface.launch()
+# Run the FastAPI server
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)