ragilbuaj's picture
change the model
b743a4c verified
raw
history blame
1.35 kB
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
import torch
# Inisialisasi model dan tokenizer
model_name = "ragilbuaj/sentiment-analysis-TWS-reviews"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
# Inisialisasi FastAPI
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Bisa disesuaikan dengan daftar asal yang diizinkan
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Model request body
class TextInput(BaseModel):
text: str
# Fungsi untuk analisis sentimen
def predict_sentiment(text):
nlp = pipeline(
"sentiment-analysis",
model=model_name,
tokenizer=model_name
)
result = nlp(text)[0]
sentiment = result['label']
confidence = result['score']
return sentiment, confidence
# Endpoint untuk analisis sentimen
@app.post("/predict")
async def predict(input: TextInput):
sentiment, confidence = predict_sentiment(input.text)
return {"sentiment": sentiment, "confidence": confidence}
# Endpoint root
@app.get("/")
async def read_root():
return {"message": "Sentiment Analysis API"}