Update api.py

api.py

@@ -1,4 +1,6 @@
-from flask import Flask, request, send_file, jsonify
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import FileResponse, JSONResponse
 from transformers import DistilBertForSequenceClassification, DistilBertTokenizer
 import pandas as pd
 import torch
@@ -6,10 +8,17 @@ import tempfile
 import os
 import re
 from collections import Counter
-from flask_cors import CORS
 
-app = Flask(__name__)
-CORS(app)
+app = FastAPI()
+
+# Enable CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 # Load model
 model = DistilBertForSequenceClassification.from_pretrained("./fine_tuned_model")
@@ -19,6 +28,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 model.eval()
 
+
 # Sentiment prediction
 def predict_sentiment(texts):
     encodings = tokenizer(texts, truncation=True, padding=True, max_length=128, return_tensors="pt")
@@ -29,6 +39,7 @@ def predict_sentiment(texts):
     sentiment_map = {0: "Negative", 1: "Neutral", 2: "Positive"}
     return [sentiment_map[p.item()] for p in predictions]
 
+
 # Frequent words
 def get_top_words(texts, n=30):
     all_words = []
@@ -39,6 +50,7 @@ def get_top_words(texts, n=30):
     most_common = counter.most_common(n)
     return pd.DataFrame(most_common, columns=['word', 'count'])
 
+
 # Identify column
 def get_text_column(df):
     for col in ['content', 'tweet', 'text']:
@@ -46,28 +58,24 @@ def get_text_column(df):
             return col
     return None
 
-# POST /predict
-@app.route('/predict', methods=['POST'])
-def predict():
-    if 'file' not in request.files:
-        return jsonify({'error': 'No file uploaded'}), 400
 
-    file = request.files['file']
+# POST /predict
+@app.post("/predict")
+async def predict(file: UploadFile = File(...)):
     try:
-        df = pd.read_csv(file)
+        df = pd.read_csv(file.file)
     except Exception:
         try:
-            file.seek(0)
-            df = pd.read_excel(file)
+            file.file.seek(0)
+            df = pd.read_excel(file.file)
         except Exception:
-            return jsonify({'error': 'Unable to read the file'}), 400
+            raise HTTPException(status_code=400, detail="Unable to read the file")
 
     text_col = get_text_column(df)
     if not text_col:
-        return jsonify({'error': 'No "content", "tweet", or "text" column found'}), 400
+        raise HTTPException(status_code=400, detail='No "content", "tweet", or "text" column found')
 
     texts = df[text_col].astype(str).tolist()
-
     df['sentiment'] = predict_sentiment(texts)
     df['content_length'] = df[text_col].astype(str).apply(len)
 
@@ -75,50 +83,47 @@ def predict():
 
     temp_dir = tempfile.mkdtemp()
     sentiment_path = os.path.join(temp_dir, 'final_data.csv')
-    df.to_csv(sentiment_path, index=False)
-
     words_path = os.path.join(temp_dir, 'word_frequent.csv')
+
+    df.to_csv(sentiment_path, index=False)
     top_words_df.to_csv(words_path, index=False)
 
-    return jsonify({
+    return JSONResponse({
         'sentiment_file': f'/download?file={sentiment_path}',
         'top_words_file': f'/download?file={words_path}',
         'sentiment_data': df.to_dict(orient='records'),
         'top_words_data': top_words_df.to_dict(orient='records')
     })
 
-# POST /wordcloud
-@app.route('/wordcloud', methods=['POST'])
-def wordcloud():
-    if 'file' not in request.files:
-        return jsonify({'error': 'No file uploaded'}), 400
 
-    file = request.files['file']
+# POST /wordcloud
+@app.post("/wordcloud")
+async def wordcloud(file: UploadFile = File(...)):
     try:
-        df = pd.read_csv(file)
+        df = pd.read_csv(file.file)
     except Exception:
         try:
-            file.seek(0)
-            df = pd.read_excel(file)
+            file.file.seek(0)
+            df = pd.read_excel(file.file)
         except Exception:
-            return jsonify({'error': 'Unable to read the file'}), 400
+            raise HTTPException(status_code=400, detail="Unable to read the file")
 
     text_col = get_text_column(df)
     if not text_col:
-        return jsonify({'error': 'No "content", "tweet", or "text" column found'}), 400
+        raise HTTPException(status_code=400, detail='No "content", "tweet", or "text" column found')
 
     texts = df[text_col].astype(str).tolist()
     top_words_df = get_top_words(texts)
 
-    return jsonify({'top_words_data': top_words_df.to_dict(orient='records')})
+    return JSONResponse({'top_words_data': top_words_df.to_dict(orient='records')})
+
 
 # GET /download
-@app.route('/download')
-def download():
-    file_path = request.args.get('file')
-    if not file_path or not os.path.exists(file_path):
-        return jsonify({'error': 'File not found'}), 404
-    return send_file(file_path, as_attachment=True)
-
-if __name__ == '__main__':
-    app.run(host="0.0.0.0", port=7860, debug=True)
+@app.get("/download")
+async def download(file: str):
+    if not file or not os.path.exists(file):
+        raise HTTPException(status_code=404, detail="File not found")
+    return FileResponse(file, filename=os.path.basename(file))
+
+
+# Run with: uvicorn main:app --host 0.0.0.0 --port 7860