Reverted changes

main.py

@@ -36,29 +36,36 @@ def format_prompt(system, message, history):
     prompt += {"role": "user", "content": message}
     return prompt
 
-def setup():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+# def setup():
+#     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    # if torch.backends.mps.is_available():
-    #     device = torch.device("mps")
-    #     x = torch.ones(1, device=device)
-    #     print (x)
-    # else:
-    #     device="cpu"
-    #     print ("MPS device not found.")
+#     # if torch.backends.mps.is_available():
+#     #     device = torch.device("mps")
+#     #     x = torch.ones(1, device=device)
+#     #     print (x)
+#     # else:
+#     #     device="cpu"
+#     #     print ("MPS device not found.")
+    
+#     # device = "auto"
+#     # device=torch.device("cpu")
     
-    # device = "auto"
-    # device=torch.device("cpu")
+#     model_path = "ibm-granite/granite-34b-code-instruct-8k"
+#     tokenizer = AutoTokenizer.from_pretrained(model_path)
+#     # drop device_map if running on CPU
+#     model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)
+#     model.eval()
     
+#     return model, tokenizer, device
+
+def generate(item: Item):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
     model_path = "ibm-granite/granite-34b-code-instruct-8k"
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     # drop device_map if running on CPU
     model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)
     model.eval()
-    
-    return model, tokenizer, device
-
-def generate(item: Item, model, tokenizer, device):
     # change input text as desired
     chat = format_prompt(item.system_prompt, item.prompt, item.history)
     chat = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
@@ -73,11 +80,12 @@ def generate(item: Item, model, tokenizer, device):
     return output_text
 
 
-model, tokenizer, device = setup()
+# model, tokenizer, device = setup()
 
 @app.post("/generate/")
 async def generate_text(item: Item):
-    return {"response": generate(item, model, tokenizer, device)}
+    return {"response": generate(item)}
+    # return {"response": generate(item, model, tokenizer, device)}
 
 @app.get("/")
 async def generate_text_root(item: Item):