Tonic commited on
Commit
fd508d7
·
verified ·
1 Parent(s): df717c1

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +93 -0
app.py ADDED
@@ -0,0 +1,93 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import spaces
2
+ import torch
3
+ import torch.nn.functional as F
4
+ from torch import Tensor
5
+ from transformers import AutoTokenizer, AutoModel
6
+ import gradio as gr
7
+
8
+ title = """
9
+ # 👋🏻Welcome to 🙋🏻‍♂️Tonic's 🐣e5-mistral🛌🏻Embeddings """
10
+ description = """
11
+ You can use this Space to test out the current model [intfloat/e5-mistral-7b-instruct](https://huggingface.co/intfloat/e5-mistral-7b-instruct). e5mistral has a larger context window, a different prompting/return mechanism and generally better results than other embedding models.
12
+ You can also use 🐣e5-mistral🛌🏻 by cloning this space. 🧬🔬🔍 Simply click here: <a style="display:inline-block" href="https://huggingface.co/spaces/Tonic/e5?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></h3>
13
+ Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder's🛠️community on 👻Discord: [![Let's build the future of AI together! 🚀🤖](https://discordapp.com/api/guilds/1109943800132010065/widget.png)](https://discord.gg/GWpVpekp) On 🤗Huggingface: [TeamTonic](https://huggingface.co/TeamTonic) & [MultiTransformer](https://huggingface.co/MultiTransformer) On 🌐Github: [Polytonic](https://github.com/tonic-ai) & contribute to 🌟 [Poly](https://github.com/tonic-ai/poly)
14
+ """
15
+ # Define the function to pool the last token
16
+ def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
17
+ left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
18
+ if left_padding:
19
+ return last_hidden_states[:, -1]
20
+ else:
21
+ sequence_lengths = attention_mask.sum(dim=1) - 1
22
+ batch_size = last_hidden_states.shape[0]
23
+ return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
24
+
25
+ # Define the function to get detailed instruct
26
+ def get_detailed_instruct(task_description: str, query: str) -> str:
27
+ return f'Instruct: {task_description}\nQuery: {query}'
28
+
29
+ # Load tokenizer and model
30
+ tokenizer = AutoTokenizer.from_pretrained('intfloat/e5-mistral-7b-instruct')
31
+ model = AutoModel.from_pretrained('intfloat/e5-mistral-7b-instruct')
32
+
33
+ @spaces.GPU
34
+ def compute_embeddings(*input_texts):
35
+ # Check if GPU is available and use it; otherwise, use CPU
36
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
37
+
38
+ # Move model to the chosen device
39
+ model.to(device)
40
+ max_length = 4096
41
+ task = 'Given a web search query, retrieve relevant passages that answer the query'
42
+
43
+ # Prepare the input texts
44
+ processed_texts = [get_detailed_instruct(task, text) for text in input_texts]
45
+
46
+ # Tokenize the input texts
47
+ batch_dict = tokenizer(processed_texts, max_length=max_length - 1, return_attention_mask=False, padding=False, truncation=True)
48
+ batch_dict['input_ids'] = [input_ids + [tokenizer.eos_token_id] for input_ids in batch_dict['input_ids']]
49
+ batch_dict = tokenizer.pad(batch_dict, padding=True, return_attention_mask=True, return_tensors='pt')
50
+
51
+ # Get model outputs
52
+ outputs = model(**batch_dict)
53
+ embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
54
+
55
+ # Normalize embeddings
56
+ embeddings = F.normalize(embeddings, p=2, dim=1)
57
+ return embeddings.detach().cpu().numpy()
58
+
59
+
60
+ def app_interface():
61
+ with gr.Blocks() as demo:
62
+ gr.Markdown(title)
63
+ gr.Markdown(description)
64
+
65
+ # Input text boxes
66
+ input_text_boxes = [gr.Textbox(label=f"Input Text {i+1}") for i in range(4)]
67
+
68
+ # Button to compute embeddings
69
+ compute_button = gr.Button("Compute Embeddings")
70
+
71
+ # Output display
72
+ output_display = gr.Dataframe(headers=["Embedding"], datatype=["numpy"])
73
+
74
+ # Layout
75
+ with gr.Row():
76
+ with gr.Column():
77
+ for text_box in input_text_boxes:
78
+ text_box.render()
79
+ with gr.Column():
80
+ compute_button.render()
81
+ output_display.render()
82
+
83
+ # Function call
84
+ compute_button.click(
85
+ fn=compute_embeddings,
86
+ inputs=input_text_boxes,
87
+ outputs=output_display
88
+ )
89
+
90
+ return demo
91
+
92
+ # Run the Gradio app
93
+ app_interface().launch()