Update app.py

reconfigure autmatic optimisations

app.py

@@ -4,6 +4,7 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import logging
 from typing import Optional, Dict, Any
 import os
+import torch
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -41,13 +42,24 @@ def load_model(model_name: str):
         logger.info(f"Loading {model_name}...")
         try:
             model_path = MODELS[model_name]
-            tokenizer = AutoTokenizer.from_pretrained(model_path, token=HF_TOKEN)
+            
+            # Load tokenizer with minimal settings
+            tokenizer = AutoTokenizer.from_pretrained(
+                model_path, 
+                token=HF_TOKEN,
+                use_fast=False  # Use slower but more stable tokenizer
+            )
+            
+            # Load model with minimal settings
             model = AutoModelForSeq2SeqLM.from_pretrained(
                 model_path,
                 token=HF_TOKEN,
-                device_map="auto",
-                torch_dtype="auto"
+                torch_dtype=torch.float32,  # Use standard precision
             )
+            
+            # Move model to CPU explicitly
+            model = model.cpu()
+            
             loaded_models[model_name] = model
             loaded_tokenizers[model_name] = tokenizer
             logger.info(f"Successfully loaded {model_name}")
@@ -89,10 +101,28 @@ async def predict(request: PredictionRequest):
         # Prepend the shared prefix
         full_input = "Interpret this dream: " + request.inputs
 
-        # Tokenize and generate
-        input_ids = tokenizer(full_input, return_tensors="pt").input_ids
-        outputs = model.generate(input_ids, **request.parameters)
-        decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Tokenize and generate with explicit error handling
+        try:
+            input_ids = tokenizer(
+                full_input, 
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=512
+            ).input_ids
+            
+            outputs = model.generate(
+                input_ids,
+                max_length=200,
+                num_return_sequences=1,
+                no_repeat_ngram_size=2,
+                **request.parameters
+            )
+            
+            decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        except Exception as e:
+            logger.error(f"Error in model prediction pipeline: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"Model prediction failed: {str(e)}")
 
         return PredictionResponse(generated_text=decoded)