Update app.py

app.py

@@ -7,20 +7,26 @@ import plotly.graph_objects as go
 
 TOKEN = os.getenv("HF_TOKEN")
 
-model_name = "mistralai/Mistral-7B-v0.3"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+default_model_name = "mistralai/Mistral-7B-Instruct-v0.3"
+tokenizer = None
 model = None
 
-# Set pad token to eos token if not defined
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-
 @spaces.GPU(duration=300)
-def get_embedding(text):
-    global model
-    if model is None:
-        model = AutoModel.from_pretrained(model_name, torch_dtype=torch.float16).cuda()
-        model.resize_token_embeddings(len(tokenizer))
+def get_embedding(text, model_repo):
+    global tokenizer, model
+    
+    if tokenizer is None or model is None or model.name_or_path != model_repo:
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(model_repo)
+            model = AutoModel.from_pretrained(model_repo, torch_dtype=torch.float16).cuda()
+            
+            # Set pad token to eos token if not defined
+            if tokenizer.pad_token is None:
+                tokenizer.pad_token = tokenizer.eos_token
+            
+            model.resize_token_embeddings(len(tokenizer))
+        except Exception as e:
+            return f"Error loading model: {str(e)}"
     
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512).to('cuda')
     with torch.no_grad():
@@ -31,8 +37,18 @@ def reduce_to_3d(embedding):
     return embedding[:3]
 
 @spaces.GPU
-def compare_embeddings(*texts):
-    embeddings = [get_embedding(text) for text in texts if text.strip()]
+def compare_embeddings(model_repo, *texts):
+    if not model_repo:
+        model_repo = default_model_name
+    
+    embeddings = []
+    for text in texts:
+        if text.strip():
+            emb = get_embedding(text, model_repo)
+            if isinstance(emb, str):  # Error message
+                return emb
+            embeddings.append(emb)
+    
     embeddings_3d = [reduce_to_3d(emb) for emb in embeddings]
     
     fig = go.Figure()
@@ -50,8 +66,9 @@ def generate_text_boxes(n):
 
 with gr.Blocks() as iface:
     gr.Markdown("# 3D Embedding Comparison")
-    gr.Markdown("Compare the embeddings of multiple strings visualized in 3D space using Llama 3.1.")
+    gr.Markdown("Compare the embeddings of multiple strings visualized in 3D space using a custom model.")
     
+    model_repo_input = gr.Textbox(label="Model Repository", value=default_model_name, placeholder="Enter the model repository (e.g., mistralai/Mistral-7B-Instruct-v0.3)")
     num_texts = gr.Slider(minimum=2, maximum=10, step=1, value=2, label="Number of texts to compare")
     
     with gr.Column() as input_column:
@@ -72,8 +89,8 @@ with gr.Blocks() as iface:
     
     compare_button.click(
         compare_embeddings,
-        inputs=text_boxes,
+        inputs=[model_repo_input] + text_boxes,
         outputs=output
     )
 
-iface.launch()
+iface.launch()