atharvsawant96's picture
Update app.py
5a9082d verified
raw
history blame contribute delete
754 Bytes
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)