Update inference.py

inference.py

@@ -6,7 +6,10 @@ from safetensors.torch import load_file
 from datasets import load_dataset
 from pprintpp import pprint
 
+import os
 
+# Đặt biến môi trường HF_HOME
+os.environ["HF_HOME"] = "/tmp/hf_cache" 
 
 class InferenceModel:
     """
@@ -14,6 +17,10 @@ class InferenceModel:
     """
 
     def __init__(self, path_to_weights, huggingface_model=True):
+
+        self.config = {
+            "hf_model_name": "deepset/roberta-base-squad2" # Ví dụ model của bạn
+        }
         ### Init Config with either Huggingface Backbone or our own ###
         self.config = RobertaConfig(pretrained_backbone="pretrained_huggingface" if huggingface_model else "random")
 
@@ -57,48 +64,4 @@ class InferenceModel:
 
         return prediction
 
-
-if __name__ == "__main__":
-
-    dataset = load_dataset("stanfordnlp/coqa")
-    
-    data = dataset["validation"][2]
-    # data = dataset["train"][0]
-    # print("answer:", data["answers"])
-    ### Sample Text ###
-    context = data["story"]
-    print("context:", context)
-    question = data["questions"][4]
-
-    tokenizer = RobertaTokenizerFast.from_pretrained("deepset/roberta-base-squad2")
-
-    encoded = tokenizer(
-        question,
-        context,
-        max_length=512,
-        truncation="only_second",
-        padding="max_length",
-        return_offsets_mapping=True,
-        return_tensors="pt"
-    )
-    offset_mapping = encoded["offset_mapping"][0].tolist()  # convert to list of tuples
-    input_ids = encoded["input_ids"][0]
-
-
-    ### Inference Model ###
-    path_to_weights = "model/RoBERTa/save_model/model.safetensors" 
-    inferencer = InferenceModel(path_to_weights=path_to_weights, huggingface_model=True)
-    prediction = inferencer.inference_model(question, context)
-    print("\n----------------------------------")
-    print("results:", prediction)
-
-    start_token_idx = prediction["start_token_idx"]
-    end_token_idx = prediction["end_token_idx"]
-
-    start_char = offset_mapping[start_token_idx][0]
-    end_char = offset_mapping[end_token_idx][1]
-    
-    print("Question:", question)
-    print("Recovered answer:", context[start_char:end_char])
-
     # test model