Update README.md

README.md

@@ -26,6 +26,19 @@ import re
 
 from transformers import pipeline
 
+prefix = """Extract SDTM domain entities from: """
+
+unstructured_text = """<ENTER YOUR TEXT HERE>"""
+
+generator = pipeline(
+    "text-generation", model="ClinicalIntelligence/saama_gemma", device="cuda"
+)
+output = generator(
+    [{"role": "user", "content": prefix + unstructured_text}],
+    return_full_text=False,
+)[0]
+llm_output = output["generated_text"]
+
 
 def extract_entities(text):
     """
@@ -55,19 +68,6 @@ def extract_entities(text):
     return extracted_data
 
 
-prefix = """Extract SDTM domain entities from: """
-
-unstructured_text = """<ENTER YOUR TEXT HERE>"""
-
-generator = pipeline(
-    "text-generation", model="ClinicalIntelligence/saama_gemma", device="cuda"
-)
-output = generator(
-    [{"role": "user", "content": prefix + unstructured_text}],
-    max_new_tokens=16000,
-    return_full_text=False,
-)[0]
-llm_output = output["generated_text"]
 extracted_entities_list = extract_entities(llm_output)
 
 for extracted_entity in extracted_entities_list: