Update app.py

app.py

@@ -1,51 +1,32 @@
 from fastapi import FastAPI, HTTPException, status, APIRouter, Request
 from pydantic import BaseModel, ValidationError
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForMaskedLM
 import torch
 import logging
-import asyncio # For running synchronous model inference in a separate thread
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 app = FastAPI(
-    title="Masked Language Model API (via TinyLlama)",
-    description="An API to perform Masked Language Modeling using a locally hosted TinyLlama model.",
+    title="NeuroBERT-Tiny Masked Language Model API",
+    description="An API to perform Masked Language Modeling using the boltuix/NeuroBERT-Tiny model.",
     version="1.0.0"
 )
 
 api_router = APIRouter()
 
-# --- TinyLlama Model Configuration ---
-# Using TinyLlama-1.1B-Chat-v1.0 which is a small, Llama-like model suitable for local inference.
-MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-# -----------------------------------
+# --- NeuroBERT-Tiny Model Configuration ---
+# Using boltuix/NeuroBERT-Tiny for Masked Language Modeling.
+MODEL_NAME = "boltuix/NeuroBERT-Tiny"
+# ----------------------------------------
 
-# Load model and tokenizer globally to avoid reloading on each request
+# Load model globally to avoid reloading on each request
 # This block runs once when the FastAPI application starts.
 try:
     logger.info(f"Loading tokenizer and model for {MODEL_NAME}...")
-    # Load tokenizer and model for Causal LM (text generation)
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    # Using torch_dtype=torch.bfloat16 for potential memory/speed benefits if GPU is available
-    # and to fit within common memory limits. Also using device_map="auto" to load efficiently.
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME, 
-        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, 
-        device_map="auto"
-    )
+    model = AutoModelForMaskedLM.from_pretrained(MODEL_NAME)
     model.eval() # Set model to evaluation mode
-    
-    # Create a text generation pipeline
-    # We will adjust this pipeline's behavior in predict_masked_lm
-    # to simulate masked LM functionality by prompting the LLM.
-    text_generator = pipeline(
-        "text-generation", 
-        model=model, 
-        tokenizer=tokenizer,
-        # Ensure pad_token_id is set if tokenizer does not have one, to avoid warnings/errors
-        pad_token_id=tokenizer.eos_token_id if tokenizer.pad_token_id is None else tokenizer.pad_token_id
-    )
     logger.info("Model loaded successfully.")
 except Exception as e:
     logger.exception(f"Failed to load model or tokenizer for {MODEL_NAME} during startup!")
@@ -61,124 +42,69 @@ class InferenceRequest(BaseModel):
 class PredictionResult(BaseModel):
     """
     Response model for individual predictions from the /predict endpoint.
-    Simplified to focus on the sequence and score, abstracting token details.
     """
     sequence: str    # The full sequence with the predicted token filled in
-    score: float     # Confidence score of the prediction (approximated for generative LLMs)
-
-async def run_inference_blocking(generator_pipeline, prompt, num_return_sequences=5):
-    """
-    Runs the synchronous model inference in a separate thread to avoid blocking FastAPI's event loop.
-    """
-    return generator_pipeline(
-        prompt,
-        max_new_tokens=10, # Generate a small number of tokens for the mask
-        num_return_sequences=num_return_sequences,
-        do_sample=True,    # Enable sampling for varied predictions
-        temperature=0.7,   # Control randomness
-        top_k=50,          # Consider top K tokens for sampling
-        top_p=0.95,        # Consider tokens up to a certain cumulative probability
-        # The stop_sequence ensures it doesn't generate too much beyond the expected word
-        stop_sequence=[" ", ".", ",", "!", "?", "\n"] # Stop after generating a word/punctuation
-    )
-
+    score: float     # Confidence score of the prediction
+    token: int       # The ID of the predicted token
+    token_str: str   # The string representation of the predicted token
 
 @api_router.post(
-    "/predict", # Prediction endpoint remains /predict
+    "/predict", # Prediction endpoint
     response_model=list[PredictionResult],
-    summary="Predicts masked tokens in a given text using a local TinyLlama model",
-    description="Accepts a text string with '[MASK]' tokens and returns up to 5 single-word predictions for each masked position using a local generative AI model. Output is simplified to sequences and scores."
+    summary="Predicts masked tokens in a given text using NeuroBERT-Tiny",
+    description="Accepts a text string with '[MASK]' tokens and returns top 5 predictions for each masked position."
 )
 async def predict_masked_lm(request: InferenceRequest):
     """
-    Predicts the most likely tokens for [MASK] positions in the input text using the TinyLlama model.
-    Returns a list of top predictions for each masked token, including the full sequence and an approximated score.
+    Predicts the most likely tokens for [MASK] positions in the input text using the NeuroBERT-Tiny model.
+    Returns a list of top 5 predictions for each masked token, including the full sequence, score, and token details.
     """
-    text = request.text
-    logger.info(f"Received prediction request for text: '{text}'")
+    try:
+        text = request.text
+        logger.info(f"Received prediction request for text: '{text}'")
 
-    if "[MASK]" not in text:
-        logger.warning("No [MASK] token found in the input text. Returning 400 Bad Request.")
-        raise HTTPException(
-            status_code=status.HTTP_400_BAD_REQUEST,
-            detail="Input text must contain at least one '[MASK]' token."
-        )
-    
-    # Find the position of the first [MASK] token to correctly prompt the LLM
-    # And to insert predictions back into the original text
-    mask_start_index = text.find("[MASK]")
-    if mask_start_index == -1: # Should already be caught above, but as a safeguard
-        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="No '[MASK]' token found in input.")
-
-    # Craft a prompt that encourages the LLM to fill the mask.
-    # The prompt guides the generative LLM to act like a fill-mask model.
-    # Example: "The quick brown fox jumps over the [MASK] dog. The word that should replace [MASK] is:"
-    # We remove "[MASK]" from the prompt for the generative model, and then
-    # prepend a guiding phrase and append the text after the mask.
-    
-    # Split text around the first [MASK]
-    parts = text.split("[MASK]", 1)
-    if len(parts) < 2: # Should not happen if [MASK] is found
-        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Error processing mask position.")
-    
-    pre_mask_text = parts[0].strip()
-    post_mask_text = parts[1].strip()
+        inputs = tokenizer(text, return_tensors="pt")
 
-    # Construct the prompt to guide TinyLlama
-    # "Fill in the blank: 'The quick brown fox jumps over the ______ dog.' Best options are:"
-    prompt = f"Complete the missing word in the following sentence. Give 5 single-word options. Sentence: '{pre_mask_text} ____ {post_mask_text}' Options:"
+        with torch.no_grad():
+            outputs = model(**inputs)
+
+        logits = outputs.logits
+        masked_token_id = tokenizer.convert_tokens_to_ids("[MASK]")
+
+        # Find all masked tokens
+        masked_token_indices = torch.where(inputs["input_ids"] == masked_token_id)[1]
+
+        if not masked_token_indices.numel():
+            logger.warning("No [MASK] token found in the input text. Returning 400 Bad Request.")
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail="Input text must contain at least one '[MASK]' token."
+            )
 
-    try:
-        # Run inference in a separate thread to not block the main event loop
-        # The model's output will be a list of dicts, e.g., [{"generated_text": "Prompt + predicted word"}]
-        raw_predictions = await run_inference_blocking(text_generator, prompt)
-        
         results = []
-        seen_words = set() # To ensure unique predictions
-
-        for i, pred_item in enumerate(raw_predictions):
-            generated_text = pred_item.get("generated_text", "")
-            
-            # Extract only the predicted word from the generated text
-            # This is heuristic and might need fine-tuning based on actual model output
-            # We look for text that comes *after* our prompt and try to extract the first word.
-            if prompt in generated_text:
-                completion_text = generated_text.split(prompt, 1)[-1].strip()
-                # Try to extract the first word if it contains spaces
-                predicted_word = completion_text.split(' ', 1)[0].strip().replace('.', '').replace(',', '')
-                # Filter out numbers, common filler words, or very short non-alpha words
-                if not predicted_word.isalpha() or len(predicted_word) < 2:
-                    continue
-                
-                # Further refine by splitting on common word separators, taking the first valid word
-                valid_words = [w for w in predicted_word.split() if w.isalpha() and len(w) > 1]
-                if not valid_words: continue
-                predicted_word = valid_words[0].lower() # Normalize to lowercase
-
-                # Ensure unique predictions
-                if predicted_word in seen_words:
-                    continue
-                seen_words.add(predicted_word)
-
-                # Construct the full sequence with the predicted word
-                full_sequence = text.replace("[MASK]", predicted_word, 1)
-                
-                # Approximate score (generative LLMs don't give scores directly for words)
-                mock_score = 0.95 - (i * 0.01) # Slightly decrease confidence for lower ranks
+        for masked_index in masked_token_indices:
+            # Get top 5 predictions for the masked token
+            top_5_logits = torch.topk(logits[0, masked_index], 5).values
+            top_5_tokens = torch.topk(logits[0, masked_index], 5).indices
+
+            for i in range(5):
+                score = torch.nn.functional.softmax(logits[0, masked_index], dim=-1)[top_5_tokens[i]].item()
+                predicted_token_id = top_5_tokens[i].item()
+                predicted_token_str = tokenizer.decode(predicted_token_id)
+
+                # Replace the [MASK] with the predicted token for the full sequence
+                temp_input_ids = inputs["input_ids"].clone()
+                temp_input_ids[0, masked_index] = predicted_token_id
+                full_sequence = tokenizer.decode(temp_input_ids[0], skip_special_tokens=True)
 
                 results.append(PredictionResult(
                     sequence=full_sequence,
-                    score=mock_score
+                    score=score,
+                    token=predicted_token_id,
+                    token_str=predicted_token_str
                 ))
-            
-            if len(results) >= 5: # Stop after getting 5 valid results
-                break
         
-        if not results:
-            logger.warning("No valid predictions could be formatted from LLM response.")
-            raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Could not extract predictions from TinyLlama output.")
-
-        logger.info(f"Successfully processed request via TinyLlama. Returning {len(results)} predictions.")
+        logger.info(f"Successfully processed request. Returning {len(results)} predictions.")
         return results
     
     except ValidationError as e:
@@ -203,7 +129,7 @@ async def predict_masked_lm(request: InferenceRequest):
 )
 async def health_check():
     logger.info("Health check endpoint accessed.")
-    return {"message": "Masked Language Model API (via TinyLlama) is running!"}
+    return {"message": "NeuroBERT-Tiny API is running!"}
 
 app.include_router(api_router)
 
@@ -214,4 +140,4 @@ async def catch_all(request: Request, path_name: str):
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")
+    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")