Tonic commited on
Commit
78a49c2
1 Parent(s): 3d3c605

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +21 -18
app.py CHANGED
@@ -51,7 +51,7 @@ class EmbeddingModel:
51
  self.model = AutoModel.from_pretrained('intfloat/e5-mistral-7b-instruct', torch_dtype=torch.float16, device_map=device)
52
 
53
  @spaces.GPU
54
- def compute_embeddings(selected_task, input_text, system_prompt):
55
  max_length = 2042
56
  task_description = tasks[selected_task]
57
  processed_texts = [f'Instruct: {task_description}\nQuery: {input_text}']
@@ -69,17 +69,22 @@ class EmbeddingModel:
69
  @spaces.GPU
70
  def compute_similarity(self, sentence1, sentence2, extra_sentence1, extra_sentence2):
71
 
72
- sentences = [sentence1, sentence2, extra_sentence1, extra_sentence2]
73
- encoded_input = self.tokenizer(sentences, padding=True, truncation=True, return_tensors='pt').to(device)
74
- with torch.no_grad():
75
- model_output = self.model(**encoded_input)
76
-
77
- embeddings = last_token_pool(model_output.last_hidden_state, encoded_input['attention_mask'])
78
- embeddings = F.normalize(embeddings, p=2, dim=1)
79
-
 
 
 
 
80
  # Compute cosine similarity
81
- similarity1 = F.cosine_similarity(embeddings[0].unsqueeze(0), embeddings[1].unsqueeze(0)).item()
82
- similarity2 = F.cosine_similarity(embeddings[2].unsqueeze(0), embeddings[3].unsqueeze(0)).item()
 
83
  return similarity1, similarity2
84
 
85
 
@@ -92,20 +97,19 @@ def app_interface():
92
 
93
  with gr.Tab("Embedding Generation"):
94
  input_text_box = gr.Textbox(label="📖Input Text")
95
- system_prompt_box = gr.Textbox(label="🤖System Prompt (Optional)")
96
  compute_button = gr.Button("Try🐣🛌🏻e5")
97
  output_display = gr.Textbox(label="🐣e5-mistral🛌🏻 Embeddings")
98
  compute_button.click(
99
  fn=EmbeddingModel.compute_embeddings,
100
- inputs=[task_dropdown, input_text_box, system_prompt_box],
101
  outputs=output_display
102
  )
103
 
104
  with gr.Tab("Sentence Similarity"):
105
- sentence1_box = gr.Textbox(label="Sentence 1")
106
- sentence2_box = gr.Textbox(label="Sentence 2")
107
- extra_sentence1_box = gr.Textbox(label="Sentence 3")
108
- extra_sentence2_box = gr.Textbox(label="Sentence 4")
109
  similarity_button = gr.Button("Compute Similarity")
110
  similarity_output = gr.Label(label="🐣e5-mistral🛌🏻 Similarity Scores")
111
  similarity_button.click(
@@ -116,7 +120,6 @@ def app_interface():
116
 
117
  with gr.Row():
118
  with gr.Column():
119
- system_prompt_box
120
  input_text_box
121
  with gr.Column():
122
  compute_button
 
51
  self.model = AutoModel.from_pretrained('intfloat/e5-mistral-7b-instruct', torch_dtype=torch.float16, device_map=device)
52
 
53
  @spaces.GPU
54
+ def compute_embeddings(selected_task, input_text):
55
  max_length = 2042
56
  task_description = tasks[selected_task]
57
  processed_texts = [f'Instruct: {task_description}\nQuery: {input_text}']
 
69
  @spaces.GPU
70
  def compute_similarity(self, sentence1, sentence2, extra_sentence1, extra_sentence2):
71
 
72
+ # Compute embeddings for each sentence
73
+ embeddings1 = compute_embeddings(self.selected_task, sentence1)
74
+ embeddings2 = compute_embeddings(self.selected_task, sentence2)
75
+ embeddings3 = compute_embeddings(self.selected_task, extra_sentence1)
76
+ embeddings4 = compute_embeddings(self.selected_task, extra_sentence2)
77
+
78
+ # Convert embeddings to tensors
79
+ embeddings_tensor1 = torch.tensor(embeddings1).to(device)
80
+ embeddings_tensor2 = torch.tensor(embeddings2).to(device)
81
+ embeddings_tensor3 = torch.tensor(embeddings3).to(device)
82
+ embeddings_tensor4 = torch.tensor(embeddings4).to(device)
83
+
84
  # Compute cosine similarity
85
+ similarity1 = F.cosine_similarity(embeddings_tensor1, embeddings_tensor2).item()
86
+ similarity2 = F.cosine_similarity(embeddings_tensor1, embeddings_tensor3).item()
87
+ similarity3 = F.cosine_similarity(embeddings_tensor1, embeddings_tensor4).item()
88
  return similarity1, similarity2
89
 
90
 
 
97
 
98
  with gr.Tab("Embedding Generation"):
99
  input_text_box = gr.Textbox(label="📖Input Text")
 
100
  compute_button = gr.Button("Try🐣🛌🏻e5")
101
  output_display = gr.Textbox(label="🐣e5-mistral🛌🏻 Embeddings")
102
  compute_button.click(
103
  fn=EmbeddingModel.compute_embeddings,
104
+ inputs=[task_dropdown, input_text_box],
105
  outputs=output_display
106
  )
107
 
108
  with gr.Tab("Sentence Similarity"):
109
+ sentence1_box = gr.Textbox(label="'Focus Sentence' - The 'Subject'")
110
+ sentence2_box = gr.Textbox(label="'Input Sentence' - 1")
111
+ extra_sentence1_box = gr.Textbox(label="'Input Sentence' - 2")
112
+ extra_sentence2_box = gr.Textbox(label="'Input Sentence' - 3")
113
  similarity_button = gr.Button("Compute Similarity")
114
  similarity_output = gr.Label(label="🐣e5-mistral🛌🏻 Similarity Scores")
115
  similarity_button.click(
 
120
 
121
  with gr.Row():
122
  with gr.Column():
 
123
  input_text_box
124
  with gr.Column():
125
  compute_button