Bijoy09 commited on
Commit
a84c0cf
1 Parent(s): 384e68d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +21 -0
app.py CHANGED
@@ -1,20 +1,37 @@
1
  from fastapi import FastAPI, HTTPException
 
2
  from pydantic import BaseModel
3
  import torch
4
  from transformers import AutoModelForSequenceClassification, AutoTokenizer
5
  import os
 
6
 
7
  app = FastAPI()
8
 
 
 
 
 
9
  # Set the cache directory for Hugging Face
10
  os.environ['TRANSFORMERS_CACHE'] = os.getenv('TRANSFORMERS_CACHE', '/app/cache')
11
 
 
 
 
 
 
 
 
 
 
12
  # Load model and tokenizer
13
  model_name = "Bijoy09/MObilebert"
14
  try:
15
  model = AutoModelForSequenceClassification.from_pretrained(model_name)
16
  tokenizer = AutoTokenizer.from_pretrained(model_name)
 
17
  except Exception as e:
 
18
  raise RuntimeError(f"Failed to load model or tokenizer: {e}")
19
 
20
  class TextRequest(BaseModel):
@@ -23,6 +40,7 @@ class TextRequest(BaseModel):
23
  @app.post("/predict/")
24
  async def predict(request: TextRequest):
25
  try:
 
26
  model.eval()
27
  inputs = tokenizer.encode_plus(
28
  request.text,
@@ -33,11 +51,14 @@ async def predict(request: TextRequest):
33
  return_attention_mask=True,
34
  return_tensors='pt'
35
  )
 
36
  with torch.no_grad():
37
  logits = model(inputs['input_ids'], attention_mask=inputs['attention_mask']).logits
 
38
  prediction = torch.argmax(logits, dim=1).item()
39
  return {"prediction": "Spam" if prediction == 1 else "Ham"}
40
  except Exception as e:
 
41
  raise HTTPException(status_code=500, detail=f"Prediction failed: {e}")
42
 
43
  @app.get("/")
 
1
  from fastapi import FastAPI, HTTPException
2
+ from fastapi.middleware.cors import CORSMiddleware
3
  from pydantic import BaseModel
4
  import torch
5
  from transformers import AutoModelForSequenceClassification, AutoTokenizer
6
  import os
7
+ import logging
8
 
9
  app = FastAPI()
10
 
11
+ # Configure logging
12
+ logging.basicConfig(level=logging.INFO)
13
+ logger = logging.getLogger(__name__)
14
+
15
  # Set the cache directory for Hugging Face
16
  os.environ['TRANSFORMERS_CACHE'] = os.getenv('TRANSFORMERS_CACHE', '/app/cache')
17
 
18
+ # Enable CORS
19
+ app.add_middleware(
20
+ CORSMiddleware,
21
+ allow_origins=["*"],
22
+ allow_credentials=True,
23
+ allow_methods=["*"],
24
+ allow_headers=["*"],
25
+ )
26
+
27
  # Load model and tokenizer
28
  model_name = "Bijoy09/MObilebert"
29
  try:
30
  model = AutoModelForSequenceClassification.from_pretrained(model_name)
31
  tokenizer = AutoTokenizer.from_pretrained(model_name)
32
+ logger.info("Model and tokenizer loaded successfully")
33
  except Exception as e:
34
+ logger.error(f"Failed to load model or tokenizer: {e}")
35
  raise RuntimeError(f"Failed to load model or tokenizer: {e}")
36
 
37
  class TextRequest(BaseModel):
 
40
  @app.post("/predict/")
41
  async def predict(request: TextRequest):
42
  try:
43
+ logger.info(f"Received text: {request.text}")
44
  model.eval()
45
  inputs = tokenizer.encode_plus(
46
  request.text,
 
51
  return_attention_mask=True,
52
  return_tensors='pt'
53
  )
54
+ logger.info(f"Tokenized inputs: {inputs}")
55
  with torch.no_grad():
56
  logits = model(inputs['input_ids'], attention_mask=inputs['attention_mask']).logits
57
+ logger.info(f"Model logits: {logits}")
58
  prediction = torch.argmax(logits, dim=1).item()
59
  return {"prediction": "Spam" if prediction == 1 else "Ham"}
60
  except Exception as e:
61
+ logger.error(f"Prediction failed: {e}")
62
  raise HTTPException(status_code=500, detail=f"Prediction failed: {e}")
63
 
64
  @app.get("/")