Update app.py

app.py

@@ -1,161 +1,136 @@
-from fastapi import FastAPI, HTTPException, Query
-from pydantic import BaseModel, Field, conlist
-import torch
-import torch.nn as nn
-import os  # Import the 'os' module
+import os
 from typing import List
 
-# --- Model Definition (same as before) ---
-class NameGenderClassifierCNN(nn.Module):
-    def __init__(self, vocab_size, embedding_dim, num_filters=64, filter_sizes=[2, 3, 4], dropout=0.5):
-        super(NameGenderClassifierCNN, self).__init__()
+from fastapi import FastAPI, HTTPException, status
+from pydantic import BaseModel
+from google import genai
+from google.genai import types
+from google.protobuf.json_format import MessageToDict  # For converting to dict
 
-        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
+app = FastAPI()
 
-        self.convs = nn.ModuleList([
-            nn.Conv1d(in_channels=embedding_dim, out_channels=num_filters, kernel_size=fs)
-            for fs in filter_sizes
-        ])
+# Load API key from environment variable
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
+if not GEMINI_API_KEY:
+    raise ValueError("The GEMINI_API_KEY environment variable is not set.")
 
-        self.fc1 = nn.Linear(len(filter_sizes) * num_filters, 100)
-        self.fc2 = nn.Linear(100, 1)
-        self.dropout = nn.Dropout(dropout)
-        self.sigmoid = nn.Sigmoid()
+genai.configure(api_key=GEMINI_API_KEY)
+client = genai.GenerativeModel('gemini-pro')  # Use a consistent model. 'gemini-pro' is better for text.
 
-    def forward(self, x):
-        x = self.embedding(x)
-        x = x.transpose(1, 2)
-        conv_outputs = []
-        for conv in self.convs:
-            conv_out = torch.relu(conv(x))
-            pool_out = torch.max_pool1d(conv_out, conv_out.shape[2])
-            conv_outputs.append(pool_out.squeeze(2))
-        x = torch.cat(conv_outputs, dim=1)
-        x = self.dropout(x)
-        x = torch.relu(self.fc1(x))
-        x = self.dropout(x)
-        x = self.fc2(x)
-        return self.sigmoid(x).squeeze()
 
+class TranslationRequest(BaseModel):
+    text: str
+    target_language: str  # Accept full language name, e.g., "Telugu", "Tamil", "Hindi"
+    source_language: str = None  # Optional:  User *might* provide the source.
 
 
-# --- Utility Function (same as before, but adapted) ---
+class TranslationResponse(BaseModel):
+    translated_text: str
+    source_language: str  # Always return the detected/used source language
+    target_language: str
 
-def tokenize_name(name, char_to_idx, max_length):
-    """Tokenizes and pads a name."""
-    name = str(name).lower()
-    tokens = [char_to_idx.get(char, char_to_idx.get(' ', 1)) for char in name]
 
-    # Pad or truncate
-    if len(tokens) < max_length:
-        tokens = tokens + [char_to_idx['<PAD>']] * (max_length - len(tokens))
-    else:
-        tokens = tokens[:max_length]
 
-    return tokens
+# --- Helper Functions ---
 
+def detect_language_and_options(text: str):
+    """Detects the language of the input text and provides translation options."""
+    contents = [
+        types.Content(
+            role="user",
+            parts=[types.Part.from_text(text=text)],
+        ),
+        types.Content(
+            role="model",
+            parts=[
+                types.Part.from_text(
+                    text="""Please identify the language of the text provided and then offer translation options as numbered choices (1-5).  Use this format:  "The text is in [Language].  Choose a language to translate to: 1. [Option 1], 2. [Option 2], 3. [Option 3], 4. [Option 4], 5. [Option 5]"."""
+                )
+            ]
+        )
+    ]
 
-# --- FastAPI Setup ---
+    response = client.generate_content(contents)
+    # Extract language and make options consistent.  Robust parsing.
+    try:
+        response_text = response.text
+        source_language = response_text.split("The text is in ")[1].split(".")[0].strip()
 
-app = FastAPI(title="Indian Name Gender Prediction API",
-              description="Predicts the gender of Indian names using a CNN model.",
-              version="1.0")
+        options_str = response_text.split("Choose a language to translate to:")[1].strip()
+        options_list = [opt.split(". ")[1].strip() for opt in options_str.split(", ")]
+           # Ensure we have *exactly* 5 options, padding if needed.
+        while len(options_list) < 5:
+            options_list.append("Option Not Available")  # Or some other placeholder
+        options_list = options_list[:5]
 
-# --- Model Loading (on startup) ---
 
-MODEL_PATH = "models/indian_name_gender_model.pt"  # Correct path within the space
+        options = {
+            str(i + 1): lang for i, lang in enumerate(options_list)
+        }
+        return source_language, options
 
+    except (IndexError, AttributeError):  # Handle parsing errors gracefully
+      raise HTTPException(status_code=500, detail="Error processing language detection response.")
 
-def load_model():
-    """Loads the model, char_to_idx, and max_name_length."""
-    try:
-        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        checkpoint = torch.load(MODEL_PATH, map_location=device)
-        char_to_idx = checkpoint['char_to_idx']
-        max_name_length = checkpoint['max_name_length']
-        config = checkpoint['model_config']
-
-        model = NameGenderClassifierCNN(
-            vocab_size=config['vocab_size'],
-            embedding_dim=config['embedding_dim'],
-            num_filters=config['num_filters'],
-            filter_sizes=config['filter_sizes']
-        )
-        model.load_state_dict(checkpoint['model_state_dict'])
-        model.to(device)
-        model.eval()  # Set to evaluation mode
-        return model, char_to_idx, max_name_length, device
-    except Exception as e:
-        raise Exception(f"Error loading model: {e}")
 
-# Load model at startup
-try:
-    model, char_to_idx, max_name_length, device = load_model()
-except Exception as e:
-    print(f"Failed to load model: {e}")
-    raise  # Re-raise the exception to halt startup
 
-# --- Pydantic Models (for request/response validation) ---
+def translate_with_gemini(text: str, source_language: str, target_language: str) -> str:
+    """Translates text using Gemini Pro, handling language codes correctly."""
 
-class PredictionRequest(BaseModel):
-    names: conlist(str, min_length=1) = Field(..., example=["Aarav", "Anika"])
-    threshold: float = Field(0.5, ge=0.0, le=1.0, description="Probability threshold for classifying as male.")
+    # More direct prompting style.  No few-shot examples needed for a simple translation task.
+    prompt = f"Translate the following text from {source_language} to {target_language}:\n\n{text}"
+    response = client.generate_content(prompt)
+
+    try:
+        return response.text
+    except (AttributeError, IndexError) as e:
+        raise HTTPException(status_code=500, detail=f"Error from Gemini API: {e}")
 
-class PredictionResponse(BaseModel):
-    predictions: List[dict] = Field(..., example=[
-        {"name": "Aarav", "predicted_gender": "Male", "male_probability": 0.95, "confidence": 0.95},
-        {"name": "Anika", "predicted_gender": "Female", "male_probability": 0.05, "confidence": 0.95}
-    ])
 
 
-# --- Prediction Function ---
 
-def predict_gender(name: str, model, char_to_idx, max_length, device, threshold: float = 0.5) -> tuple[str, float, float]:
-    """Predicts gender for a single name.  Includes threshold."""
-    tokenized_name = tokenize_name(name, char_to_idx, max_length)
-    input_tensor = torch.tensor([tokenized_name], dtype=torch.long).to(device)
+@app.post("/translate", response_model=TranslationResponse, status_code=status.HTTP_200_OK)
+async def translate(request: TranslationRequest):
+    """Translates text from a source language to a target language."""
 
-    with torch.no_grad():
-        output = model(input_tensor)
-        probability = output.item()
-        predicted_gender = 'Male' if probability >= threshold else 'Female'
-        confidence = probability if probability >= threshold else 1 - probability
-    return predicted_gender, probability, confidence
+    if not request.text:
+        raise HTTPException(status_code=400, detail="Text to translate cannot be empty.")
+    if not request.target_language:
+        raise HTTPException(status_code=400, detail="Target language must be provided.")
 
-# --- API Endpoints ---
 
-@app.get("/", response_model=str)
-async def read_root():
-	return "Welcome to the Indian Name Gender Prediction API.  Use the /predict endpoint."
+    if request.source_language:  # User provided source language.  Use it directly.
+         source_language = request.source_language
 
-@app.post("/predict", response_model=PredictionResponse)
-async def predict(request: PredictionRequest):
-    """Predicts the gender of one or more Indian names."""
+    else:  # Detect the language
+        try:
+            source_language, _ = detect_language_and_options(request.text)  # We don't need options here
+        except HTTPException as e: # Re-raise HTTP exceptions from the helper function.
+            raise e
+        except Exception as e:  # Catch any other unexpected errors.
+            raise HTTPException(status_code=500, detail=f"Language detection failed: {e}")
+    # Validate the target language against a reasonable set of supported languages.
+    supported_languages = ["English", "Hindi", "Telugu", "Marathi", "Bengali", "Tamil", "Spanish", "French", "German", "Japanese", "Chinese"]  # Add more as needed.
+    if request.target_language not in supported_languages:
+      raise HTTPException(status_code=400, detail=f"Target language '{request.target_language}' is not supported. Supported languages: {', '.join(supported_languages)}")
     try:
-        predictions = []
-        for name in request.names:
-            gender, prob, conf = predict_gender(name, model, char_to_idx, max_name_length, device, request.threshold)
-            predictions.append({
-                "name": name,
-                "predicted_gender": gender,
-                "male_probability": prob,
-                "confidence": conf
-            })
-        return {"predictions": predictions}
+        translated_text = translate_with_gemini(request.text, source_language, request.target_language)
+        return TranslationResponse(translated_text=translated_text, source_language=source_language, target_language=request.target_language)
+    except HTTPException as e:
+        raise e
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-    
-@app.get("/predict_single")
-async def predict_single(name: str = Query(..., description="The name to predict."),
-                         threshold: float = Query(0.5, ge=0.0, le=1.0, description="Probability threshold for classifying as male.")):
-    """Predicts gender for a *single* name, provided as a query parameter."""
+        raise HTTPException(status_code=500, detail=f"Translation failed: {e}")
+
+
+@app.post("/detect_language", status_code=status.HTTP_200_OK)
+async def detect_language(text: str = ""):  # Simpler input, just the text
+    """Detects the language of the input text and provides translation options."""
+    if not text:
+      raise HTTPException(status_code=400, detail="Text to detect cannot be empty.")
     try:
-        gender, prob, conf = predict_gender(name, model, char_to_idx, max_name_length, device, threshold)
-        return {
-            "name": name,
-            "predicted_gender": gender,
-            "male_probability": prob,
-            "confidence": conf
-        }
+      source_language, options = detect_language_and_options(text)
+      return {"source_language": source_language, "translation_options": options}
+    except HTTPException as e:
+      raise e
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+      raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}")