Update app.py

app.py

@@ -1,29 +1,24 @@
-from fastapi import FastAPI, Request
-from sentence_transformers import SentenceTransformer
+import gradio as gr
+from transformers import AutoTokenizer, AutoModel
 import torch
 
-# Initialize FastAPI
-app = FastAPI()
-
-# Load your SentenceTransformer model with trust_remote_code=True
+# Load your model with trust_remote_code=True
 model = SentenceTransformer("ICTuniverse/tuned-bi-encoder", trust_remote_code=True)
 
-# Use GPU if available
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-
-@app.post("/embed")
-async def embed_text(request: Request):
-    data = await request.json()
-    text = data.get("text")
-    if not text:
-        return {"error": "Text is required"}
-    
-    # Get embeddings
+# Define a function to get embeddings
+def get_embedding(text):
     with torch.no_grad():
-        embeddings = model.encode([text], device=device, convert_to_tensor=True)
-    
-    # Convert embeddings to list for JSON serialization
-    embedding_list = embeddings.cpu().tolist()
-    
-    return {"embedding": embedding_list}
+        embeddings = model.encode(text)
+    return embeddings
+
+# Create a Gradio interface
+iface = gr.Interface(
+    fn=get_embedding,
+    inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
+    outputs="json",
+    title="Embedding Generator",
+    description="Get embeddings using ICTuniverse/tuned-bi-encoder"
+)
+
+# Launch the Gradio app
+iface.launch()