piepeline

metadata_transformer.py

@@ -8,17 +8,19 @@ TOKEN = os.environ.get("TOKEN")
 
 login(token=TOKEN)
 
-model = "meta-llama/Llama-2-7b-chat-hf"
+def translate(schema_input, schema_target):
 
-tokenizer = AutoTokenizer.from_pretrained(model)
-pipeline = transformers.pipeline(
-    "text-generation",
-    model=model,
-    torch_dtype=torch.float16,
-    device_map="auto",
-)
 
-def translate(schema_input, schema_target, pipeline):
+    model = "meta-llama/Llama-2-7b-chat-hf"
+
+    tokenizer = AutoTokenizer.from_pretrained(model)
+    pipeline = transformers.pipeline(
+        "text-generation",
+        model=model,
+        torch_dtype=torch.float16,
+        device_map="auto",
+    )
+
     sequences = pipeline(
         '{} \n Translate the schema metadata file above to the schema: {}'.format(schema_input, schema_target),
         do_sample=True,