v2p1

Beta. Revert if trash.

app.py

@@ -4,7 +4,8 @@ import torch
 from transformers import AutoTokenizer, AutoModel
 import plotly.graph_objects as go
 
-model_name = "mistralai/Mistral-7B-v0.1"
+# Update the model name to Llama 3.1
+model_name = "meta-llama/Meta-Llama-3.1-405B-Instruct-FP8"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = None
 
@@ -16,7 +17,7 @@ if tokenizer.pad_token is None:
 def get_embedding(text):
     global model
     if model is None:
-        model = AutoModel.from_pretrained(model_name).cuda()
+        model = AutoModel.from_pretrained(model_name, torch_dtype=torch.float16).cuda()
         model.resize_token_embeddings(len(tokenizer))
     
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512).to('cuda')
@@ -47,7 +48,7 @@ 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 Mistral 7B.")
+    gr.Markdown("Compare the embeddings of multiple strings visualized in 3D space using Llama 3.1.")
     
     num_texts = gr.Slider(minimum=2, maximum=10, step=1, value=2, label="Number of texts to compare")
     
@@ -73,4 +74,4 @@ with gr.Blocks() as iface:
         outputs=output
     )
 
-iface.launch()
+iface.launch()