Update app.py

app.py

@@ -2,16 +2,17 @@ import gradio as gr
 import torch
 import faiss
 import numpy as np
-import pandas as pd
 import datasets
 from transformers import AutoTokenizer, AutoModel
 
 title = "HouseMD bot"
 
-description = "Gradio Demo for telegram bot \
-To use it, simply add your text message. \
+description = "Gradio Demo for telegram bot.\
+To use it, simply add your text message.\n\
 I've used the API on this Space to deploy the model on a Telegram bot."
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 
 def embed_bert_cls(text, model, tokenizer):
     t = tokenizer(text, padding=True, truncation=True, return_tensors='pt')
@@ -29,23 +30,20 @@ def get_ranked_docs(query, vec_query_base, data,
     index = faiss.IndexFlatL2(vec_shape)
     index.add(vec_query_base)
     xq = embed_bert_cls(query, bi_model, bi_tok)
-    D, I = index.search(xq.reshape(1, vec_shape), 50)
-
-    corpus = []
-    for i in I[0]:
-        corpus.append(data['answer'][i])
+    _, I = index.search(xq.reshape(1, vec_shape), 50)
+    corpus = [data[int(i)]['answer'] for i in I[0]]
 
     queries = [query] * len(corpus)
     tokenized_texts = cross_tok(
         queries, corpus, max_length=128, padding=True, truncation=True, return_tensors="pt"
-    ).to(config.model.device)
+    ).to(device)
 
     with torch.no_grad():
-        ce_scores = cross_model(
-            tokenized_texts['input_ids'], tokenized_texts['attention_mask']
-        ).last_hidden_state[:, 0, :]
-        ce_scores = ce_scores @ ce_scores.T
-
+        model_output = cross_model(
+            **{k: v.to(cross_model.device) for k, v in tokenized_texts.items()}
+        )
+    ce_scores = model_output.last_hidden_state[:, 0, :]
+    ce_scores = np.matmul(ce_scores, ce_scores.T)
     scores = ce_scores.cpu().numpy()
     scores_ix = np.argsort(scores)[::-1]
 
@@ -55,18 +53,15 @@ def get_ranked_docs(query, vec_query_base, data,
 def load_dataset(url='ekaterinatao/house_md_context3'):
 
     dataset = datasets.load_dataset(url, split='train')
-    house_dataset = []
-    for data in dataset:
-        if data['labels'] == 0:
-            house_dataset.append(data)
+    house_dataset = dataset.filter(lambda row: row['labels'] == 0)
 
-    return pd.DataFrame(house_dataset)
+    return house_dataset
 
 
 def load_cls_base(url='ekaterinatao/house_md_cls_embeds'):
 
     cls_dataset = datasets.load_dataset(url, split='train')
-    cls_base = np.stack([embed for embed in pd.DataFrame(cls_dataset)['cls_embeds']])
+    cls_base = np.stack([embed['cls_embeds'] for embed in cls_dataset])
 
     return cls_base
 
@@ -101,6 +96,7 @@ def get_answer(message):
     )
     return answer
 
+
 interface = gr.Interface(
     fn=get_answer,
     inputs=gr.Textbox(label="Input message to House MD", lines=3),