Update app.py

app.py

@@ -1,31 +1,49 @@
-import gradio as gr
-import torch
-import pandas as pd
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 from transformers import MarianMTModel, MarianTokenizer
+import torch
+import time
 import logging
-import os
 
 # ---------------- LOGGING ----------------
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
-logger = logging.getLogger("FrenchTitleApp")
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger("FrenchTitleAPI")
+
+# ---------------- FASTAPI APP ----------------
+app = FastAPI(title="French Title Generator API")
 
 # ---------------- MODEL CLASS ----------------
 class FrenchTitleGenerator:
-    def __init__(self, model_name="Helsinki-NLP/opus-mt-en-fr"):
-        logger.info(f"Loading model: {model_name}")
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.tokenizer = MarianTokenizer.from_pretrained(model_name)
-        self.model = MarianMTModel.from_pretrained(model_name).to(self.device)
-        logger.info("Model loaded successfully")
+    def __init__(self, model_name='Helsinki-NLP/opus-mt-en-fr'):
+        try:
+            logger.info("Loading tokenizer and model...")
+            self.tokenizer = MarianTokenizer.from_pretrained(model_name)
+            self.model = MarianMTModel.from_pretrained(model_name)
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            self.model.to(self.device)
+            self.model.eval()
+            logger.info(f"Model loaded on {self.device}")
+        except Exception as e:
+            logger.error(f"Model initialization failed: {e}")
+            raise
 
-    def generate(self, year, city, original_title, max_length=50, temperature=0.8):
+    def generate_title(self, year: str, city: str, original_title: str, max_length: int = 50, temperature: float = 0.8):
         try:
             input_text = f"Generate title from: Year: {year} City: {city} Description: {original_title}"
-            inputs = self.tokenizer(input_text, return_tensors="pt", padding=True, truncation=True).to(self.device)
+            inputs = self.tokenizer(
+                input_text,
+                return_tensors="pt",
+                max_length=128,
+                truncation=True,
+                padding=True
+            ).to(self.device)
 
             with torch.no_grad():
                 outputs = self.model.generate(
-                    inputs["input_ids"],
+                    input_ids=inputs["input_ids"],
                     attention_mask=inputs["attention_mask"],
                     max_length=max_length,
                     temperature=temperature,
@@ -33,38 +51,55 @@ class FrenchTitleGenerator:
                     num_beams=3,
                     early_stopping=True
                 )
-
-            result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            for phrase in ["Le titre de", "Année:", "Ville:", "Description:", "[YEAR]", "[CITY]", "[DESC]", "[TITLE]"]:
-                result = result.replace(phrase, "")
-            return " ".join(result.split()).strip()
-
+            return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
         except Exception as e:
-            logger.error(f"Error during generation: {e}")
-            return f"Error: {e}"
+            logger.error(f"Title generation failed: {e}")
+            return f"{original_title} - {city} {year}"
 
-# ---------------- INIT MODEL ----------------
-generator = FrenchTitleGenerator()
+# ---------------- INITIALIZE MODEL ----------------
+model = FrenchTitleGenerator()
 
-# ---------------- GRADIO APP ----------------
-def generate_title(year, city, original_title, temperature=0.8, max_length=50):
-    return generator.generate(year, city, original_title, max_length=max_length, temperature=temperature)
+# ---------------- REQUEST SCHEMA ----------------
+class TitleRequest(BaseModel):
+    year: str
+    city: str
+    original_title: str
+    max_length: int = 50
+    temperature: float = 0.8
 
-iface = gr.Interface(
-    fn=generate_title,
-    inputs=[
-        gr.Textbox(label="Year", placeholder="e.g. 2025"),
-        gr.Textbox(label="City", placeholder="e.g. Paris"),
-        gr.Textbox(label="Original Title", placeholder="e.g. Jazz Festival in the City Park"),
-        gr.Slider(0.1, 1.5, value=0.8, step=0.1, label="Temperature"),
-        gr.Slider(10, 100, value=50, step=5, label="Max Length")
-    ],
-    outputs=gr.Textbox(label="Generated French Title"),
-    title="🇫🇷 French Title Generator",
-    description="Generate a French event title based on year, city, and description using a translation model.",
-    theme="soft"
-)
+# ---------------- API ENDPOINTS ----------------
+@app.get("/api/health")
+def health_check():
+    return {
+        "status": "healthy",
+        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+        "model_loaded": model is not None
+    }
+
+@app.post("/api/generate")
+def generate_title(req: TitleRequest):
+    try:
+        result = model.generate_title(
+            year=req.year,
+            city=req.city,
+            original_title=req.original_title,
+            max_length=req.max_length,
+            temperature=req.temperature
+        )
+        return {
+            "input": req.dict(),
+            "generated_title": result,
+            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
-if __name__ == "__main__":
-    iface.launch()
-    
+@app.get("/")
+def root():
+    return {
+        "message": "French Title Generator API is running",
+        "endpoints": {
+            "/api/health": "Health check",
+            "/api/generate": "POST JSON to generate title"
+        }
+    }