kells1986 commited on
Commit
b38ebdd
·
1 Parent(s): bd097d0

Added an app and requirements file

Browse files
Files changed (2) hide show
  1. app.py +59 -0
  2. requirements.txt +2 -0
app.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoTokenizer, AutoModel
3
+ import torch
4
+ import torch.nn.functional as F
5
+
6
+
7
+ # Mean Pooling - Take attention mask into account for correct averaging
8
+ def mean_pooling(model_output, attention_mask):
9
+ token_embeddings = model_output[0] # First element of model_output contains all token embeddings
10
+ input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
11
+ return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
12
+
13
+
14
+ class Matcher:
15
+
16
+ def __init__(self):
17
+ self.tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
18
+ self.model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
19
+
20
+ def _encoder(self, text: list[str]):
21
+ encoded_input = self.tokenizer(text, padding=True, truncation=True, return_tensors='pt')
22
+ with torch.no_grad():
23
+ model_output = self.model(**encoded_input)
24
+ sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
25
+ sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
26
+ return sentence_embeddings
27
+
28
+ def __call__(self, textA: list[str], textB: list[str]):
29
+ embeddings_a = self._encoder(textA)
30
+ embeddings_b = self._encoder(textB)
31
+ sim = embeddings_a @ embeddings_b.T
32
+ match_inds = torch.argmax(sim, dim=1)
33
+ match_conf = torch.max(sim, dim=1).values
34
+ return match_inds.tolist(), match_conf.tolist()
35
+
36
+
37
+ def run_match(source_text, destination_text):
38
+ matcher = Matcher()
39
+ sources = source_text.split("\n")
40
+ destinations = destination_text.split("\n")
41
+ match_inds, match_conf = matcher(sources, destinations)
42
+ matches = [f"{sources[i]} -> {destinations[match_inds[i]]} ({match_conf[i]:.2f})" for i in
43
+ range(len(sources))]
44
+ return "\n".join(matches)
45
+
46
+
47
+ with gr.Blocks() as demo:
48
+ with gr.Row():
49
+ with gr.Column():
50
+ source_text = gr.Textbox(lines=10, label="Source Text", name="source_text")
51
+ with gr.Column():
52
+ dest_text = gr.Textbox(lines=10, label="Destination Text", name="destination_text")
53
+ with gr.Column():
54
+ matches = gr.Textbox(lines=10, label="Matches", name="matches")
55
+ with gr.Row():
56
+ match_btn = gr.Button(label="Match", name="run")
57
+ match_btn.click(fn=run_match, inputs=[source_text, dest_text], outputs=matches)
58
+
59
+ demo.launch()
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ torch==2.0.0
2
+ transformers==4.25.1