Spaces:
Sleeping
Sleeping
| from fastapi import FastAPI, HTTPException | |
| from pydantic import BaseModel | |
| from clinical_ner import ClinicalNER | |
| import uvicorn | |
| # Initialize FastAPI app | |
| app = FastAPI( | |
| title="Clinical NER API", | |
| description="Named Entity Recognition API using Bio_ClinicalBERT", | |
| version="1.0.0" | |
| ) | |
| # Initialize the NER model (singleton pattern) | |
| ner_model = None | |
| async def startup_event(): | |
| """Load the NER model on startup""" | |
| global ner_model | |
| ner_model = ClinicalNER() | |
| print("NER model loaded successfully!") | |
| # Request model | |
| class TextRequest(BaseModel): | |
| text: str | |
| class Config: | |
| json_schema_extra = { | |
| "example": { | |
| "text": "Patient presents with hypertension and diabetes. Prescribed metformin 500mg." | |
| } | |
| } | |
| # Response models | |
| class Entity(BaseModel): | |
| entity_group: str | |
| score: float | |
| word: str | |
| start: int | |
| end: int | |
| class BasicNERResponse(BaseModel): | |
| entities: list[Entity] | |
| count: int | |
| class PrologNERResponse(BaseModel): | |
| prolog_facts: str | |
| count: int | |
| async def root(): | |
| """Root endpoint with API information""" | |
| return { | |
| "message": "Clinical NER API", | |
| "endpoints": { | |
| "/ner/basic": "POST - Get basic NER annotations", | |
| "/ner/prolog": "POST - Get Prolog facts", | |
| "/docs": "GET - Interactive API documentation" | |
| } | |
| } | |
| async def ner_basic(request: TextRequest): | |
| """ | |
| Perform basic NER on the input text. | |
| Returns a list of detected entities with their types, positions, and confidence scores. | |
| """ | |
| try: | |
| if not request.text.strip(): | |
| raise HTTPException(status_code=400, detail="Text cannot be empty") | |
| entities = ner_model.basic_ner(request.text) | |
| return { | |
| "entities": entities, | |
| "count": len(entities) | |
| } | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Error processing text: {str(e)}") | |
| async def ner_prolog(request: TextRequest): | |
| """ | |
| Perform NER and return results as Prolog facts. | |
| Returns Prolog facts in the format: entity(Id, Type, Word, Start, End, Score). | |
| """ | |
| try: | |
| if not request.text.strip(): | |
| raise HTTPException(status_code=400, detail="Text cannot be empty") | |
| prolog_facts = ner_model.prolog_ner(request.text) | |
| # Count the number of facts (lines) | |
| count = len(prolog_facts.split('\n')) if prolog_facts else 0 | |
| return { | |
| "prolog_facts": prolog_facts, | |
| "count": count | |
| } | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Error processing text: {str(e)}") | |
| async def health_check(): | |
| """Health check endpoint""" | |
| return { | |
| "status": "healthy", | |
| "model_loaded": ner_model is not None | |
| } | |
| if __name__ == "__main__": | |
| uvicorn.run(app, host="0.0.0.0", port=8000) | |