only run from local files when running locally

.idea/workspace.xml

@@ -6,7 +6,10 @@
   <component name="ChangeListManager">
     <list default="true" id="ba1eab7b-54ae-409d-bac6-f23bdc636598" name="Changes" comment="">
       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/app.py" beforeDir="false" afterPath="$PROJECT_DIR$/app.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/ai.py" beforeDir="false" afterPath="$PROJECT_DIR$/ai.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/cpu_ai.py" beforeDir="false" afterPath="$PROJECT_DIR$/cpu_ai.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/models/abacaj/replit-v2-codeinstruct-3b.q4_1.bin" beforeDir="false" />
+      <change beforePath="$PROJECT_DIR$/models/marella/ggml-model.bin" beforeDir="false" />
     </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />

ai.py

@@ -3,8 +3,8 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import cpu_ai
 
 models = [
-    "replit-v2-codeinstruct-3b.q4_1", #"abacaj/Replit-v2-CodeInstruct-3B-ggml",
-    "ggml-model", #"marella/gpt-2-ggml",
+    "abacaj/Replit-v2-CodeInstruct-3B-ggml",
+    "marella/gpt-2-ggml",
     "WizardLM/WizardCoder-Python-34B-V1.0",
     "WizardLM/WizardCoder-15B-V1.0",
     "WizardLM/WizardCoder-Python-7B-V1.0",
@@ -45,9 +45,9 @@ def cleanup_response(generated_text):
 def generate_code(prompt, model_index, max_tokens, temperature=0.6):
     model_full_name = models[model_index]
     if model_index == 0:
-        output = cpu_ai.generate_code(prompt, "models/abacaj", model_full_name, max_tokens, temperature)
+        output = cpu_ai.generate_code(prompt, model_full_name, max_tokens, temperature)
     elif model_index == 1:
-        output = cpu_ai.generate_code(prompt, "models/marella", model_full_name, max_tokens, temperature)
+        output = cpu_ai.generate_code(prompt, model_full_name, max_tokens, temperature)
     else:
         output = run_general_model(model_full_name, prompt, max_tokens, temperature)
 

cpu_ai.py

@@ -1,4 +1,5 @@
 import os
+from pathlib import Path
 from dataclasses import dataclass, asdict
 from ctransformers import AutoModelForCausalLM, AutoConfig
 
@@ -29,14 +30,21 @@ def generate(llm: AutoModelForCausalLM,
     return llm(format_prompt(prompt), **asdict(generation_config))
 
 
-def generate_code(prompt, config_path, model_name, max_tokens, temperature):
-    path = os.path.abspath(f"{config_path}/{model_name}.bin")
+def generate_code(prompt, model_name, max_tokens, temperature):
+    from_local = False
+
+    model_path = model_name
+    config_path = model_name
+    if from_local:
+        config_folder = model_name.split("/")[0]
+        config_path = os.path.abspath(f"models/{config_folder}")
+        model_path = os.path.abspath(f"models/{model_name}.bin")
 
     config = AutoConfig.from_pretrained(
-        os.path.abspath(config_path),
+        config_path,
     )
     llm = AutoModelForCausalLM.from_pretrained(
-        path,
+        model_path,
         model_type="replit",
         config=config,
     )