from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import os

# Set the cache directory to a writable location
os.environ['TRANSFORMERS_CACHE'] = '/app/.cache'
os.environ['HF_HOME'] = '/app/.cache'

# Inisialisasi model dan tokenizer
model_name = "w11wo/indonesian-roberta-base-sentiment-classifier"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Inisialisasi FastAPI
app = FastAPI()

# Model request body
class TextInput(BaseModel):
    text: str

# Fungsi untuk analisis sentimen
def predict_sentiment(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    outputs = model(**inputs)
    scores = outputs.logits[0].detach().numpy()
    predictions = torch.nn.functional.softmax(torch.tensor(scores), dim=0)
    sentiment = torch.argmax(predictions).item()
    return sentiment, predictions[sentiment].item()

# 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"}