Convert to Gradio interface with API compatibility

app.py

@@ -1,9 +1,7 @@
-from flask import Flask, request, jsonify
+import gradio as gr
 from transformers import AutoTokenizer, AutoModel
 import torch
-import os
-
-app = Flask(__name__)
+import json
 
 # Load your model once
 model_name = "sentence-transformers/all-MiniLM-L6-v2"
@@ -21,16 +19,46 @@ def get_embedding(text):
         embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
     return embeddings.squeeze().tolist()
 
-@app.route('/api/predict', methods=['POST'])
-def predict():
+def predict_texts(texts):
+    """Generate embeddings for a list of texts (for API compatibility)"""
+    if isinstance(texts, str):
+        # If single text, convert to list
+        texts = [texts]
+
+    if not isinstance(texts, list):
+        return "Error: Input must be a list of texts or a single text string"
+
+    # Generate embeddings for each text
+    embeddings = []
+    for text in texts:
+        if isinstance(text, str):
+            embedding = get_embedding(text)
+            embeddings.append(embedding)
+        else:
+            return f"Error: All items must be strings, got {type(text)}"
+
+    return embeddings
+
+def predict_single_text(text):
+    """Generate embedding for a single text (for Gradio interface)"""
+    if not text or not text.strip():
+        return "Please enter some text to generate embeddings."
+
+    embedding = get_embedding(text.strip())
+    return f"Embedding (first 10 values): {embedding[:10]}...\nFull embedding has {len(embedding)} dimensions."
+
+def predict_api(json_str):
+    """Handle API calls from backend - expects JSON string with {"data": ["text1", "text2", ...]}"""
     try:
-        data = request.get_json()
+        import json
+        data = json.loads(json_str)
+
         if not data or 'data' not in data:
-            return jsonify({'error': 'Missing data field'}), 400
+            return json.dumps({'error': 'Missing data field'})
 
         texts = data['data']
         if not isinstance(texts, list):
-            return jsonify({'error': 'Data must be a list of texts'}), 400
+            return json.dumps({'error': 'Data must be a list of texts'})
 
         # Generate embeddings for each text
         embeddings = []
@@ -39,17 +67,41 @@ def predict():
                 embedding = get_embedding(text)
                 embeddings.append(embedding)
             else:
-                return jsonify({'error': 'All items in data must be strings'}), 400
-
-        return jsonify({'data': embeddings})
+                return json.dumps({'error': 'All items must be strings'})
 
+        return json.dumps({'data': embeddings})
     except Exception as e:
-        return jsonify({'error': str(e)}), 500
+        return json.dumps({'error': str(e)})
 
-@app.route('/health', methods=['GET'])
-def health():
-    return jsonify({'status': 'healthy'})
+# Create API interface (this will create /api/predict endpoint)
+api_interface = gr.Interface(
+    fn=predict_api,
+    inputs=gr.Textbox(),
+    outputs=gr.Textbox(),
+    api_name="predict"
+)
+
+# Create web interface
+web_interface = gr.Interface(
+    fn=predict_single_text,
+    inputs=gr.Textbox(lines=3, placeholder="Enter text to generate embeddings..."),
+    outputs=gr.Textbox(label="Embedding Result"),
+    title="Text Embedding Generator",
+    description="Generate embeddings for text using sentence-transformers/all-MiniLM-L6-v2 model",
+    examples=[
+        ["Hello world"],
+        ["This is a test sentence for embedding generation."],
+        ["Machine learning is transforming the world."]
+    ]
+)
+
+# Launch both interfaces
+if __name__ == '__main__':
+    gr.TabbedInterface([web_interface, api_interface], ["Web UI", "API"]).launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True
+    )
 
 if __name__ == '__main__':
-    port = int(os.environ.get('PORT', 7860))
-    app.run(host='0.0.0.0', port=port)
+    iface.launch(server_name="0.0.0.0", server_port=7860, share=True)

requirements.txt

@@ -1,4 +1,4 @@
-Flask==2.3.3
+gradio==4.36.0
 transformers==4.36.0
 torch>=2.0.0
 numpy>=1.21.0