Update app.py

app.py

@@ -1,14 +1,56 @@
 import gradio as gr
 import spaces
 import torch
+from transformers import AutoTokenizer, AutoModel
+from sklearn.decomposition import PCA
+import plotly.graph_objects as go
+from huggingface_hub import HfApi
+from huggingface_hub import hf_hub_download
+import os
+import sys
 
 zero = torch.Tensor([0]).cuda()
 print(zero.device) # <-- 'cpu' 🤔
 
+model_name = "sentence-transformers/all-MiniLM-L6-v2"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModel.from_pretrained(model_name)
+
+@spaces.GPU
+def get_embedding(text):
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    return outputs.last_hidden_state.mean(dim=1).squeeze().numpy()
+
+def compress_to_3d(embedding):
+    pca = PCA(n_components=3)
+    return pca.fit_transform(embedding.reshape(1, -1))[0]
+
 @spaces.GPU
-def greet(n):
-    print(zero.device) # <-- 'cuda:0' 🤗
-    return f"Hello {zero + n} Tensor"
+def compare_embeddings(text1, text2):
+    emb1 = get_embedding(text1)
+    emb2 = get_embedding(text2)
+    
+    emb1_3d = compress_to_3d(emb1)
+    emb2_3d = compress_to_3d(emb2)
+    
+    fig = go.Figure(data=[
+        go.Scatter3d(x=[0, emb1_3d[0]], y=[0, emb1_3d[1]], z=[0, emb1_3d[2]], mode='lines+markers', name='Text 1'),
+        go.Scatter3d(x=[0, emb2_3d[0]], y=[0, emb2_3d[1]], z=[0, emb2_3d[2]], mode='lines+markers', name='Text 2')
+    ])
+    
+    fig.update_layout(scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'))
+    
+    return fig
 
-demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
-demo.launch()
+iface = gr.Interface(
+    fn=compare_embeddings,
+    inputs=[
+        gr.Textbox(label="Text 1"),
+        gr.Textbox(label="Text 2")
+    ],
+    outputs=gr.Plot(),
+    title="3D Embedding Comparison",
+    description="Compare the embeddings of two strings visualized in 3D space."
+)