Update app.py

app.py

@@ -148,7 +148,8 @@ theme = gr.themes.Monochrome(
 
 def stream(model, code, generate_kwargs):
     input_ids = tokenizer(code, return_tensors="pt").to(device)
-    generated_ids = model.generate(**input_ids, **generate_kwargs)
+    # generated_ids = model.generate(**input_ids, **generate_kwargs)
+    generated_ids = model.generate(**input_ids)
     return tokenizer.decode(generated_ids[0][input_ids["input_ids"].shape[1]:], skip_special_tokens=True).strip()
 
 @spaces.GPU
@@ -170,23 +171,22 @@ def generate(
         seed=42,
     )
 
-    # if method == "Base":
-    #     output = stream(basemodel, prompt, generate_kwargs)
-    # elif method == "Prefix":
-    #     output = stream(model_map[library + " Prefix"], prompt, generate_kwargs)
-    # elif method == "Evo Prefix" and library in ["SQLModel", "SfePy", "MegEngine"]:
-    #     output = stream(model_map["Main Evo Prefix"], prompt, generate_kwargs)
-    # elif method == "FFT" and library in ["SQLModel", "SfePy", "MegEngine"]:
-    #     output = stream(model_map[library + " FFT"], prompt, generate_kwargs)
-    # elif method == "Evo FFT" and library in ["SQLModel", "SfePy", "MegEngine"]:
-    #     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 = ""
-    output = ""
+    if method == "Base":
+        output = stream(basemodel, prompt, generate_kwargs)
+    elif method == "Prefix":
+        output = stream(model_map[library + " Prefix"], prompt, generate_kwargs)
+    elif method == "Evo Prefix" and library in ["SQLModel", "SfePy", "MegEngine"]:
+        output = stream(model_map["Main Evo Prefix"], prompt, generate_kwargs)
+    elif method == "FFT" and library in ["SQLModel", "SfePy", "MegEngine"]:
+        output = stream(model_map[library + " FFT"], prompt, generate_kwargs)
+    elif method == "Evo FFT" and library in ["SQLModel", "SfePy", "MegEngine"]:
+        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 = ""
     return output