Update README.md

README.md

@@ -75,7 +75,72 @@ Users (both direct and downstream) should be made aware of the risks, biases and
 
 Use the code below to get started with the model.
 
-[More Information Needed]
+[
+import torch
+from transformers import pipeline
+import os
+import json
+
+class GeoLLMBertInference:
+    def __init__(self, config_path='config.json'):
+        with open(config_path, 'r') as config_file:
+            config = json.load(config_file)
+        
+        self.project_path = config['project_path']
+        self.tokenizer_path = os.path.join(self.project_path, config['tokenizer_path'])
+        self.model_path = os.path.join(self.project_path, config['model_path'])
+
+        # Check if a GPU is available and set the device accordingly
+        self.device = 0 if torch.cuda.is_available() else -1
+ 
+        self.ner_pipeline = pipeline("ner", model=self.model_path, tokenizer=self.tokenizer_path, device=self.device)
+        self.result = None
+        self.concatenate_result = None
+
+    def get_ner_result(self, address):
+        self.result = self.ner_pipeline(address.upper())
+        return self.result
+
+    def concatenate_entities(self):
+        if self.result is None:
+            raise ValueError("NER result is not available. Please run get_ner_result first.")
+        
+        concatenated_result = {}
+        for entity in self.result:
+            tag = entity['entity']
+            word = entity['word'].replace('##', '').replace(',', '')
+            if tag not in concatenated_result:
+                concatenated_result[tag] = word.upper()
+            else:
+                concatenated_result[tag] += '' + word.upper()
+        
+        self.concatenate_result = concatenated_result
+        return self.concatenate_result
+
+    def get_json_result(self):
+        if self.concatenate_result is None:
+            raise ValueError("Concatenated result is not available. Please run concatenate_entities first.")
+        
+        return json.dumps(self.concatenate_result, indent=4)
+
+# Example Usage
+if __name__ == "__main__":
+    geo_llm = GeoLLMBertInference('code/geo_llm/config.json')
+    address = "16 ChSeAStREtST.CATHARINE"
+    result = geo_llm.get_ner_result(address)
+    print(result)
+
+    concatenate_result = geo_llm.concatenate_entities()
+    print(concatenate_result)
+
+    # Get the concatenated result in JSON format
+    json_result = geo_llm.get_json_result()
+    data = json.loads(json_result)
+
+
+    # Print the JSON string
+    print(json_result)
+]
 
 ## Training Details