from fastapi import FastAPI
from pydantic import BaseModel
from transformers import RobertaTokenizer, RobertaForMaskedLM, pipeline

app = FastAPI()

# Load SmartBERT v3
model_name = "web3se/SmartBERT-v3"
tokenizer = RobertaTokenizer.from_pretrained(model_name)
model = RobertaForMaskedLM.from_pretrained(model_name)

# Define API input format
class ContractRequest(BaseModel):
    contract_code: str

@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} 

# Run the FastAPI server
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)