Update app.py

app.py

@@ -7,7 +7,6 @@ import torch.nn as nn
 import transformers
 from transformers import AutoTokenizer, AutoConfig, LlamaForCausalLM, LlamaTokenizer, GenerationConfig
 
-
 auth_token = os.environ.get("AUTH_TOKEN_SECRET")
 
 tokenizer = LlamaTokenizer.from_pretrained("Claimed/capybara", use_auth_token=auth_token)
@@ -19,7 +18,44 @@ model = LlamaForCausalLM.from_pretrained(
 
 #model = model.to('cuda')
 
-
+def broad_scope_class_predictor(class_embeddings, abstract_embedding, N=5, Sensitivity='Medium'):
+    """
+    Takes in pre-computed class embeddings and abstract texts, converts abstract text into
+    :param class_embeddings: dataframe of class embeddings
+    :param abstract: a single abstract embedding
+    :param N: N highest matching classes to return, from highest to lowest, default is 5
+    :return: predictions: a full dataframe of all the predictions on the 9500+ classes, HighestSimilarity: Dataframe of the N most similar classes
+    """
+    predictions = pd.DataFrame(columns=['Class Name', 'Score'])
+    for i in range(len(class_embeddings)):
+        class_name = class_embeddings.iloc[i, 0]
+        embedding = class_embeddings.iloc[i, 2]
+        embedding = convert_saved_embeddings(embedding)
+        abstract_embedding = abstract_embedding.numpy()
+        abstract_embedding = torch.from_numpy(abstract_embedding)
+        cos = torch.nn.CosineSimilarity(dim=1)
+        score = cos(abstract_embedding, embedding).numpy().tolist()
+        result = [class_name, score[0]]
+        predictions.loc[len(predictions)] = result
+    greenpredictions = predictions.tail(52)
+    if Sensitivity == 'High':
+        Threshold = 0.5
+    elif Sensitivity == 'Medium':
+        Threshold = 0.40
+    elif Sensitivity == 'Low':
+        Threshold = 0.35
+    GreenLikelihood = 'False'
+    for i in range(len(greenpredictions)):
+        score = greenpredictions.iloc[i, 1]
+        if float(score) >= Threshold:
+            GreenLikelihood = 'True'
+            break
+        else:
+            continue
+    HighestSimilarity = predictions.nlargest(N, ['Score'])
+
+    return HighestSimilarity
+    
 def add_text(history, text):
     history = history + [(text, None)]
     return history, ""
@@ -42,12 +78,11 @@ def classifier(userin):
     in_emb = classification.sentence_embedder(clean_in, 'Model_bert')
 
     Number = 10
-    broad_scope_predictions = classification.broad_scope_class_predictor(class_embeddings, in_emb, Number, Sensitivity='High')
+    broad_scope_predictions = broad_scope_class_predictor(class_embeddings, in_emb, Number, Sensitivity='High')
 
-    return broad_scope_predictions[1]
+    return broad_scope_predictions
     
 
-
 def generateresponse(history):#, task):
     """
     Model definition here: