inference api

inference_hf/__init__.py

@@ -0,0 +1 @@
+from ._inference import InferenceHF

inference_hf/_inference.py

@@ -0,0 +1,29 @@
+import json
+import requests
+from typing import Union,List
+
+class InferenceHF:
+    headers = {"Authorization": f"Bearer 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"}
+    API_URL = "https://api-inference.huggingface.co/models/"
+
+    @classmethod
+    def inference(cls, inputs: Union[List[str], str], model_name:str) ->dict:
+        payload = dict(
+            inputs = inputs,
+            options = dict(
+                wait_for_model=True
+            )
+        )
+
+        data = json.dumps(payload)
+        response = requests.request("POST", cls.API_URL+model_name, headers=cls.headers, data=data)
+        return json.loads(response.content.decode("utf-8"))
+
+
+if __name__ == '__main__':
+    print(InferenceHF.inference(
+        inputs='hi how are you?',
+        model_name= 't5-small'
+    ))
+
+

lrt/clustering/clustering_pipeline.py

@@ -5,6 +5,7 @@ from ..utils import __create_model__
 from sklearn.cluster import KMeans
 # from yellowbrick.cluster import KElbowVisualizer
 from .clusters import ClusterList
+
 class ClusterPipeline:
     def __init__(self, config:Configuration = None):
         if config is None:

lrt/lrt.py

@@ -72,7 +72,7 @@ class LiteratureResearchTool:
                              best_k: int = 5
                              ) -> (ClusterList,ArticleList):
 
-        @st.cache(hash_funcs={Tokenizer: Tokenizer.__hash__})
+        @st.cache(hash_funcs={Tokenizer: Tokenizer.__hash__},allow_output_mutation=True)
         def ieee_process(
                 query: str,
                 num_papers: int,
@@ -87,7 +87,7 @@ class LiteratureResearchTool:
             clusters = self.__postprocess_clusters__(clusters)
             return clusters, articles
 
-        @st.cache(hash_funcs={Tokenizer: Tokenizer.__hash__})
+        @st.cache(hash_funcs={Tokenizer: Tokenizer.__hash__},allow_output_mutation=True)
         def arxiv_process(
                 query: str,
                 num_papers: int,
@@ -100,7 +100,7 @@ class LiteratureResearchTool:
             clusters = self.__postprocess_clusters__(clusters)
             return clusters, articles
 
-        @st.cache(hash_funcs={Tokenizer: Tokenizer.__hash__})
+        @st.cache(hash_funcs={Tokenizer: Tokenizer.__hash__},allow_output_mutation=True)
         def pwc_process(
                 query: str,
                 num_papers: int,

lrt/utils/functions.py

@@ -4,6 +4,7 @@ from kmeans_pytorch import kmeans
 import torch
 from sklearn.cluster import KMeans
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,Text2TextGenerationPipeline
+from inference_hf import InferenceHF
 
 class Template:
     def __init__(self):
@@ -62,64 +63,86 @@ def __create_model__(model_ckpt):
         return ret
 
     elif model_ckpt == 'keyphrase-transformer':
-        tokenizer = AutoTokenizer.from_pretrained(template.keywords_extraction[model_ckpt])
-        model = AutoModelForSeq2SeqLM.from_pretrained(template.keywords_extraction[model_ckpt])
-        pipe = Text2TextGenerationPipeline(model=model, tokenizer=tokenizer)
-
-        def ret(texts: List[str]):
-            tmp = pipe(texts)
-            results = [
-                set(
-                    map(str.strip,
-                        x['generated_text'].split('|') #[str...]
-                        )
-                )
-                for x in tmp] # [{str...}...]
-
-            return results
-
-        return ret
-
-    elif model_ckpt == 'KeyBartAdapter':
         model_ckpt = template.keywords_extraction[model_ckpt]
-        tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_ckpt)
-        pipe = Text2TextGenerationPipeline(model=model,tokenizer=tokenizer)
 
         def ret(texts: List[str]):
-            tmp = pipe(texts)
-            results = [
-                set(
-                    map(str.strip,
-                        x['generated_text'].split(';')  # [str...]
-                        )
-                )
-                for x in tmp]  # [{str...}...]
-
-            return results
+            # first try inference API
+            response = InferenceHF.inference(
+                inputs=texts,
+                model_name=model_ckpt
+            )
+
+            # inference failed:
+            if not isinstance(response, list):
+                tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
+                model = AutoModelForSeq2SeqLM.from_pretrained(model_ckpt)
+                pipe = Text2TextGenerationPipeline(model=model, tokenizer=tokenizer)
+
+                tmp = pipe(texts)
+                results = [
+                    set(
+                        map(str.strip,
+                            x['generated_text'].split('|')  # [str...]
+                            )
+                    )
+                    for x in tmp]  # [{str...}...]
+
+                return results
+
+            # inference sucsess
+            else:
+                results = [
+                    set(
+                        map(str.strip,
+                            x['generated_text'].split('|')  # [str...]
+                            )
+                    )
+                    for x in response]  # [{str...}...]
+
+                return results
 
         return ret
 
-    elif model_ckpt == 'KeyBart':
+    elif model_ckpt == 'KeyBartAdapter' or model_ckpt == 'KeyBart':
         model_ckpt = template.keywords_extraction[model_ckpt]
-        tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_ckpt)
-        pipe = Text2TextGenerationPipeline(model=model,tokenizer=tokenizer)
-
         def ret(texts: List[str]):
-            tmp = pipe(texts)
-            results = [
-                set(
-                    map(str.strip,
-                        x['generated_text'].split(';')  # [str...]
-                        )
-                )
-                for x in tmp]  # [{str...}...]
-
-            return results
+            # first try inference API
+            response = InferenceHF.inference(
+                inputs=texts,
+                model_name=model_ckpt
+            )
+
+            # inference failed:
+            if not isinstance(response,list):
+                tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
+                model = AutoModelForSeq2SeqLM.from_pretrained(model_ckpt)
+                pipe = Text2TextGenerationPipeline(model=model, tokenizer=tokenizer)
+
+
+                tmp = pipe(texts)
+                results = [
+                    set(
+                        map(str.strip,
+                            x['generated_text'].split(';')  # [str...]
+                            )
+                    )
+                    for x in tmp]  # [{str...}...]
+
+                return results
+
+            # inference sucsess
+            else:
+                results = [
+                    set(
+                        map(str.strip,
+                            x['generated_text'].split(';')  # [str...]
+                            )
+                    )
+                    for x in response]  # [{str...}...]
+
+                return results
 
         return ret
-
     else:
         raise RuntimeError(f'The model {model_ckpt} is not supported. Please open an issue on the GitHub about the model.')