Update app.py

app.py

@@ -65,7 +65,7 @@ def load_model():
             return True
         except Exception as fallback_error:
             logger.error(f"Fallback model loading also failed: {str(fallback_error)}")
-            return False
+        return False
 
 def generate_embeddings(texts: Union[str, List[str]]) -> Union[List[float], List[List[float]]]:
     """Generate embeddings for input text(s) using Qwen3-Embedding-0.6B model"""
@@ -129,18 +129,30 @@ def generate_embeddings(texts: Union[str, List[str]]) -> Union[List[float], List
                         
             except Exception as e:
                 logger.warning(f"Error generating embedding for text: {str(e)}")
-                # Return zero vector as last resort
-                embeddings.append([0.0] * 1024)  # Qwen3-Embedding-0.6B has 1024 dimensions
+                # Return zero vector as last resort - use correct dimension based on model type
+                if hasattr(model, 'config') and hasattr(model.config, 'hidden_size'):
+                    # Qwen3 model dimension
+                    embeddings.append([0.0] * model.config.hidden_size)
+                else:
+                    # Fallback model dimension (384 for all-MiniLM-L6-v2)
+                    embeddings.append([0.0] * 384)
         
         return embeddings[0] if single_text else embeddings
         
     except Exception as e:
         logger.error(f"Error in generate_embeddings: {str(e)}")
-        # Return zero vectors as fallback
+        # Return zero vectors as fallback - use correct dimension
+        if hasattr(model, 'config') and hasattr(model.config, 'hidden_size'):
+            # Qwen3 model dimension
+            fallback_dim = model.config.hidden_size
+        else:
+            # Fallback model dimension (384 for all-MiniLM-L6-v2)
+            fallback_dim = 384
+            
         if single_text:
-            return [0.0] * 1024
+            return [0.0] * fallback_dim
         else:
-            return [[0.0] * 1024] * len(texts)
+            return [[0.0] * fallback_dim] * len(texts)
 
 def compute_similarity(embedding1: List[float], embedding2: List[float]) -> float:
     """Compute cosine similarity between two embeddings"""
@@ -278,29 +290,61 @@ async def health():
 async def predict(data: dict):
     """Main prediction endpoint for embeddings"""
     try:
-        if "data" not in data:
-            raise HTTPException(status_code=400, detail="Missing 'data' field in request")
-        
-        input_data = data["data"]
+        # Check for new format first (texts parameter)
+        if "texts" in data:
+            texts = data["texts"]
+            normalize = data.get("normalize", True)
+            
+            if not isinstance(texts, list):
+                raise HTTPException(status_code=400, detail="'texts' must be a list")
+            
+            if len(texts) == 0:
+                raise HTTPException(status_code=400, detail="'texts' list cannot be empty")
+            
+            # Generate embeddings
+            logger.info(f"Generating embeddings for {len(texts)} texts")
+            embeddings = generate_embeddings(texts)
+            logger.info(f"Generated {len(embeddings)} embeddings with dimension {len(embeddings[0]) if embeddings else 0}")
+            
+            # Normalize embeddings if requested
+            if normalize:
+                import numpy as np
+                embeddings = [emb / np.linalg.norm(emb) for emb in embeddings]
+                logger.info("Embeddings normalized")
+            
+            return {
+                "embeddings": embeddings,
+                "model": MODEL_NAME,
+                "usage": {
+                    "prompt_tokens": sum(len(text.split()) for text in texts),
+                    "total_tokens": sum(len(text.split()) for text in texts)
+                }
+            }
         
-        # Handle single text or batch texts
-        if isinstance(input_data, str):
-            # Single text
-            embeddings = generate_embeddings(input_data)
-            return {"data": [embeddings]}
-        elif isinstance(input_data, list):
-            if len(input_data) > 0 and isinstance(input_data[0], str):
-                # Single text in list
-                embeddings = generate_embeddings(input_data[0])
-                return {"data": [embeddings]}
-            elif len(input_data) > 0 and isinstance(input_data[0], list):
-                # Batch texts
-                embeddings = generate_embeddings(input_data[0])
+        # Fallback to old format for backward compatibility
+        elif "data" in data:
+            input_data = data["data"]
+            
+            # Handle single text or batch texts
+            if isinstance(input_data, str):
+                # Single text
+                embeddings = generate_embeddings(input_data)
                 return {"data": [embeddings]}
+            elif isinstance(input_data, list):
+                if len(input_data) > 0 and isinstance(input_data[0], str):
+                    # Single text in list
+                    embeddings = generate_embeddings(input_data[0])
+                    return {"data": [embeddings]}
+                elif len(input_data) > 0 and isinstance(input_data[0], list):
+                    # Batch texts
+                    embeddings = generate_embeddings(input_data[0])
+                    return {"data": [embeddings]}
+                else:
+                    raise HTTPException(status_code=400, detail="Invalid data format")
             else:
-                raise HTTPException(status_code=400, detail="Invalid data format")
+                raise HTTPException(status_code=400, detail="Invalid data type")
         else:
-            raise HTTPException(status_code=400, detail="Invalid data type")
+            raise HTTPException(status_code=400, detail="Missing 'texts' or 'data' field in request")
             
     except Exception as e:
         logger.error(f"Error in predict endpoint: {str(e)}")
@@ -308,19 +352,42 @@ async def predict(data: dict):
 
 @app.post("/api/similarity")
 async def similarity(data: dict):
-    """Compute similarity between two embeddings"""
+    """Compute similarity between two texts or embeddings"""
     try:
-        if "embedding1" not in data or "embedding2" not in data:
-            raise HTTPException(status_code=400, detail="Missing embedding1 or embedding2 field")
-        
-        emb1 = data["embedding1"]
-        emb2 = data["embedding2"]
+        # Check for new format first (text1, text2 parameters)
+        if "text1" in data and "text2" in data:
+            text1 = data["text1"]
+            text2 = data["text2"]
+            
+            if not isinstance(text1, str) or not isinstance(text2, str):
+                raise HTTPException(status_code=400, detail="text1 and text2 must be strings")
+            
+            # Generate embeddings for both texts
+            emb1 = generate_embeddings(text1)
+            emb2 = generate_embeddings(text2)
+            
+            # Compute similarity
+            sim = compute_similarity(emb1, emb2)
+            return {
+                "similarity": sim,
+                "model": MODEL_NAME,
+                "text1": text1,
+                "text2": text2
+            }
         
-        if not isinstance(emb1, list) or not isinstance(emb2, list):
-            raise HTTPException(status_code=400, detail="Embeddings must be lists")
+        # Fallback to old format (embedding1, embedding2 parameters)
+        elif "embedding1" in data and "embedding2" in data:
+            emb1 = data["embedding1"]
+            emb2 = data["embedding2"]
+            
+            if not isinstance(emb1, list) or not isinstance(emb2, list):
+                raise HTTPException(status_code=400, detail="Embeddings must be lists")
+            
+            sim = compute_similarity(emb1, emb2)
+            return {"similarity": sim}
         
-        sim = compute_similarity(emb1, emb2)
-        return {"similarity": sim}
+        else:
+            raise HTTPException(status_code=400, detail="Missing 'text1' and 'text2' or 'embedding1' and 'embedding2' fields")
         
     except Exception as e:
         logger.error(f"Error in similarity endpoint: {str(e)}")