Update app.py

app.py

@@ -34,23 +34,23 @@ DSPY_PREFIX_URL = "luna-code/dspy-codegen-350M-mono-prefix"
 CS_EVO_PREFIX_URL = "luna-code/cs-codegen-350M-mono-evo-prefix"
 
 tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_URL)
-basemodel = AutoModelForCausalLM.from_pretrained(CHECKPOINT_URL)
+basemodel = AutoModelForCausalLM.from_pretrained(CHECKPOINT_URL, device_map="auto")
 
-sql_prefix = PeftModel.from_pretrained(basemodel, SQLMODEL_PREFIX_URL)
-sfepy_prefix = PeftModel.from_pretrained(basemodel, SFEPY_PREFIX_URL)
-megengine_prefix = PeftModel.from_pretrained(basemodel, MEGENGINE_PREFIX_URL)
-main_evo_prefix = PeftModel.from_pretrained(basemodel, MAIN_EVO_PREFIX_URL)
+sql_prefix = PeftModel.from_pretrained(basemodel, SQLMODEL_PREFIX_URL, device_map="auto")
+sfepy_prefix = PeftModel.from_pretrained(basemodel, SFEPY_PREFIX_URL, device_map="auto")
+megengine_prefix = PeftModel.from_pretrained(basemodel, MEGENGINE_PREFIX_URL, device_map="auto")
+main_evo_prefix = PeftModel.from_pretrained(basemodel, MAIN_EVO_PREFIX_URL, device_map="auto")
 
-sqlmodel_fft = AutoModelForCausalLM.from_pretrained(SQLMODEL_FFT_URL)
-sfepy_fft = AutoModelForCausalLM.from_pretrained(SFEPY_FFT_URL)
-megengine_fft = AutoModelForCausalLM.from_pretrained(MEGENGINE_FFT_URL)
-main_evo_fft = AutoModelForCausalLM.from_pretrained(MAIN_EVO_FFT_URL)
-main_fd_fft = AutoModelForCausalLM.from_pretrained(MAIN_FD_FFT_URL)
+sqlmodel_fft = AutoModelForCausalLM.from_pretrained(SQLMODEL_FFT_URL, device_map="auto")
+sfepy_fft = AutoModelForCausalLM.from_pretrained(SFEPY_FFT_URL, device_map="auto")
+megengine_fft = AutoModelForCausalLM.from_pretrained(MEGENGINE_FFT_URL, device_map="auto")
+main_evo_fft = AutoModelForCausalLM.from_pretrained(MAIN_EVO_FFT_URL, device_map="auto")
+main_fd_fft = AutoModelForCausalLM.from_pretrained(MAIN_FD_FFT_URL, device_map="auto")
 
-langchain_prefix = PeftModel.from_pretrained(basemodel, LANGCHAIN_PREFIX_URL)
-llamaindex_prefix = PeftModel.from_pretrained(basemodel, LLAMAINDEX_PREFIX_URL)
-dspy_prefix = PeftModel.from_pretrained(basemodel, DSPY_PREFIX_URL)
-cs_evo_prefix = PeftModel.from_pretrained(basemodel, CS_EVO_PREFIX_URL)
+langchain_prefix = PeftModel.from_pretrained(basemodel, LANGCHAIN_PREFIX_URL, device_map="auto")
+llamaindex_prefix = PeftModel.from_pretrained(basemodel, LLAMAINDEX_PREFIX_URL, device_map="auto")
+dspy_prefix = PeftModel.from_pretrained(basemodel, DSPY_PREFIX_URL, device_map="auto")
+cs_evo_prefix = PeftModel.from_pretrained(basemodel, CS_EVO_PREFIX_URL, device_map="auto")
 
 # basemodel = ""
 # sql_prefix = ""
@@ -147,8 +147,7 @@ theme = gr.themes.Monochrome(
 )
 
 def stream(model, code, generate_kwargs):
-    model.to(device)
-    input_ids = tokenizer(code, return_tensors="pt").to("cuda")
+    input_ids = tokenizer(code, return_tensors="pt").to(device)
     generated_ids = model.generate(**input_ids, **generate_kwargs)
     return tokenizer.decode(generated_ids[0][input_ids["input_ids"].shape[1]:], skip_special_tokens=True).strip()
 
@@ -183,6 +182,8 @@ def generate(
         output = stream(model_map["Main Evo FFT"], prompt, generate_kwargs)
     elif method == "Full Data FFT" and library in ["SQLModel", "SfePy", "MegEngine"]:
         output = stream(model_map["Main FD FFT"], prompt, generate_kwargs)
+    elif method == "Evo Prefix" and library in ["LangChain", "LlamaIndex", "DSPy"]:
+        output = stream(model_map["CS Evo Prefix"], prompt, generate_kwargs)
     else:
         output = ""
 
@@ -241,7 +242,7 @@ with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
         gr.Markdown(description)
         with gr.Row():
             library = gr.Dropdown(
-                        ["SQLModel", "SfePy", "MegEngine", "LangChain", "LlamaIndex", "DSpy"],
+                        ["SQLModel", "SfePy", "MegEngine", "LangChain", "LlamaIndex", "DSPy"],
                         value="LangChain",
                         label="Library",
                         info="Choose a library from the list",