Update app.py

app.py

@@ -1,98 +1,195 @@
 from fastapi import FastAPI, HTTPException, status, APIRouter, Request
 from pydantic import BaseModel, ValidationError
-from transformers import AutoTokenizer, AutoModelForMaskedLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 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="NeuroBERT-Tiny Masked Language Model API",
-    description="An API to perform Masked Language Modeling using the boltuix/NeuroBERT-Tiny model.",
+    title="Masked Language Model API (via TinyLlama)",
+    description="An API to perform Masked Language Modeling using a locally hosted TinyLlama 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"
+# -----------------------------------
+
+# Load model and tokenizer globally to avoid reloading on each request
+# This block runs once when the FastAPI application starts.
 try:
-    logger.info("Loading tokenizer and model for boltuix/NeuroBERT-Tiny...")
-    tokenizer = AutoTokenizer.from_pretrained("boltuix/NeuroBERT-Tiny")
-    model = AutoModelForMaskedLM.from_pretrained("boltuix/NeuroBERT-Tiny")
-    model.eval()
+    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.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("Failed to load model or tokenizer during startup!")
+    logger.exception(f"Failed to load model or tokenizer for {MODEL_NAME} during startup!")
     raise RuntimeError(f"Could not load model: {e}")
 
 class InferenceRequest(BaseModel):
+    """
+    Request model for the /predict endpoint.
+    Expects a single string field 'text' containing the sentence with [MASK] tokens.
+    """
     text: str
 
 class PredictionResult(BaseModel):
-    sequence: str
-    score: float
-    token: int
-    token_str: str
+    """
+    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
+    )
+
 
 @api_router.post(
-    "/predict", # Prediction endpoint
+    "/predict", # Prediction endpoint remains /predict
     response_model=list[PredictionResult],
-    summary="Predicts masked tokens in a given text",
-    description="Accepts a text string with '[MASK]' tokens and returns top 5 predictions for each masked position."
+    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."
 )
 async def predict_masked_lm(request: InferenceRequest):
-    try:
-        text = request.text
-        logger.info(f"Received prediction request for text: '{text}'")
-
-        inputs = tokenizer(text, return_tensors="pt")
-        with torch.no_grad():
-            outputs = model(**inputs)
-
-        logits = outputs.logits
-        masked_token_id = tokenizer.convert_tokens_to_ids("[MASK]")
-
-        masked_token_indices = torch.where(inputs["input_ids"] == masked_token_id)[1]
+    """
+    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.
+    """
+    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()
+
+    # 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:"
 
-        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 = []
-        for masked_index in masked_token_indices:
-            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)
+        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)
                 
-                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)
+                # Approximate score (generative LLMs don't give scores directly for words)
+                mock_score = 0.95 - (i * 0.01) # Slightly decrease confidence for lower ranks
 
                 results.append(PredictionResult(
                     sequence=full_sequence,
-                    score=score,
-                    token=predicted_token_id,
-                    token_str=predicted_token_str
+                    score=mock_score
                 ))
             
-            logger.info(f"Successfully processed request. Returning {len(results)} predictions.")
-            return results
+            if len(results) >= 5: # Stop after getting 5 valid results
+                break
         
-    except ValidationError as e: # This is line 88
+        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.")
+        return results
+    
+    except ValidationError as e:
         logger.error(f"Validation error for request: {e.errors()}")
         raise HTTPException(
             status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
             detail=e.errors()
         )
     except HTTPException:
-        raise
+        raise # Re-raise custom HTTPExceptions
     except Exception as e:
         logger.exception(f"An unexpected error occurred during prediction: {e}")
         raise HTTPException(
@@ -107,7 +204,7 @@ async def predict_masked_lm(request: InferenceRequest):
 )
 async def health_check():
     logger.info("Health check endpoint accessed.")
-    return {"message": "NeuroBERT-Tiny API is running!"}
+    return {"message": "Masked Language Model API (via TinyLlama) is running!"}
 
 app.include_router(api_router)
 
@@ -119,3 +216,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")
+