From 3b75f66eacb94e96ce5dda909d415722e628d20b Mon Sep 17 00:00:00 2001
From: peytontolbert <peytonct@gmail.com>
Date: Tue, 30 Jan 2024 16:53:41 -0500
Subject: [PATCH] print statement

---
 vision_datasets/functioncallgenerate2.py |  7 ++++---
 vision_datasets/functioncallgenerate3.py | 26 +++++++++++++++++-------
 2 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/vision_datasets/functioncallgenerate2.py b/vision_datasets/functioncallgenerate2.py
index 8c8967c..aa1c2ec 100644
--- a/vision_datasets/functioncallgenerate2.py
+++ b/vision_datasets/functioncallgenerate2.py
@@ -1,5 +1,6 @@
 import torch
 import json
+from accelerate import PartialState
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers.generation import GenerationConfig
 import re
@@ -19,11 +20,11 @@
 )
 model = AutoModelForCausalLM.from_pretrained(
     model_name_or_path,
-    device_map=("cuda:0"),
+    device_map=("cuda"),
     torch_dtype=torch.float16,
     trust_remote_code=True,
+    use_safetensors=True
 )
-
 # File to store the responses
 functions_file = "functions.json"
 
@@ -101,7 +102,7 @@ def expand_qa(features):
         )  # Extract everything after "assistant"
     else:
         response = "No response found after 'assistant'."
-
+    print(response)
     return response
 
 
diff --git a/vision_datasets/functioncallgenerate3.py b/vision_datasets/functioncallgenerate3.py
index 1135ad6..ca03ebf 100644
--- a/vision_datasets/functioncallgenerate3.py
+++ b/vision_datasets/functioncallgenerate3.py
@@ -1,5 +1,7 @@
 import torch
 import json
+import torch.distributed as dist
+import torch.multiprocessing as mp
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers.generation import GenerationConfig
 import re
@@ -19,16 +21,19 @@
 )
 model = AutoModelForCausalLM.from_pretrained(
     model_name_or_path,
-    device_map=("auto"),
+    device_map=("cuda:0"),
     torch_dtype=torch.float16,
     trust_remote_code=True,
+    use_safetensors=True
 )
 
 # File to store the responses
 functions_file = "functions.json"
+rank = 2
+world_size = 2
 
-
-def expand_qa(features):
+def expand_qa(features, rank, world_size):
+    dist.init_process_group("nccl", rank=rank, world_size=world_size)
     prompt = f"""{features}"""
     system_message = """When presented with features described by a visual language model, synthesize a function call and generate its output. The function call should be structured to capture specific attributes of the image as detailed by the visual description. Start the function call with the <fn_call> tag and then provide the expected output in JSON format.
 
@@ -81,7 +86,6 @@ def expand_qa(features):
     input_ids = tokenizer(
         prompt_template, return_tensors="pt"
     ).input_ids.cuda()
-    model.to(device)
     outputs = model.generate(
         input_ids,
         temperature=0.7,
@@ -143,9 +147,17 @@ def process_responses(file_path, output_file_path):
             save_response(item)
     return data
 
+def run_inference(rank, world_size):
+    # Process the responses.json file
+    updated_data = process_responses("responses.json", "functions.json")
 
-# Process the responses.json file
-updated_data = process_responses("responses.json", "functions.json")
+print("Data saved in functions.json")
 
 
-print("Data saved in functions.json")
+def main():
+    world_size = 2
+    mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)
+
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file