Update app.py

app.py

@@ -1,27 +1,49 @@
-from fastapi import FastAPI, UploadFile, File
-import shutil
 import os
-from image_caption import generate_caption
+
+os.environ["HF_HOME"] = "/tmp/huggingface"
+os.makedirs("/tmp/huggingface", exist_ok=True)
+from fastapi import FastAPI
+from pydantic import BaseModel
+import torch
+import numpy as np
+from transformers import AutoTokenizer, AutoModel
+from sklearn.linear_model import LogisticRegression
+import uvicorn
 
 app = FastAPI()
 
-UPLOAD_DIR = "uploads"
-os.makedirs(UPLOAD_DIR, exist_ok=True)
+# Load Hugging Face model
+model_name = "bert-base-uncased"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModel.from_pretrained(model_name)
+
+# Function to get text embeddings
+def get_embedding(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    return outputs.last_hidden_state[:, 0, :].numpy()
+
+# Sample dataset
+texts = ["I love this!", "This is terrible.", "Fantastic experience!", "I hate it.", "Absolutely wonderful!", "Worst ever!"]
+labels = [1, 0, 1, 0, 1, 0]  # 1 = Positive, 0 = Negative
+X = np.vstack([get_embedding(text) for text in texts])
+y = np.array(labels)
 
-@app.post("/upload-image/")
-async def upload_image(file: UploadFile = File(...)):
-    """Endpoint to accept image and return generated caption"""
-    file_path = f"{UPLOAD_DIR}/{file.filename}"
-    
-    # Save the uploaded file
-    with open(file_path, "wb") as buffer:
-        shutil.copyfileobj(file.file, buffer)
+# Train model
+clf = LogisticRegression()
+clf.fit(X, y)
 
-    # Generate caption
-    caption = generate_caption(file_path)
+# Define request format
+class InputText(BaseModel):
+    text: str
 
-    return {"filename": file.filename, "caption": caption}
+@app.post("/predict")
+def predict_sentiment(data: InputText):
+    user_embedding = get_embedding(data.text)
+    prediction = clf.predict(user_embedding)
+    sentiment = "Positive 😊" if prediction[0] == 1 else "Negative 😡"
+    return {"sentiment": sentiment}
 
-@app.get("/")
-async def root():
-    return {"message": "Image-to-Text API is running. Use /upload-image to send an image."}
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)