07/01/23-15:27

app.py

@@ -4,7 +4,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
 
-def extract_code(input_text):
+def extract_code_codegen(input_text):
     pattern = r"'''py\n(.*?)'''"
     match = re.search(pattern, input_text, re.DOTALL)
     
@@ -12,18 +12,27 @@ def extract_code(input_text):
         return match.group(1)
     else:
         return None  # Return None if no match is found
+    
+def extract_code_mistral(input_text):
+    pattern = r'\[CODE\](.*?)\[/CODE\]'
+    match = re.search(pattern, input_text, re.DOTALL)
+
+    if match:
+        return match.group(1)
+    else:
+        return None  # Return None if no match is found
 
 def generate_code(input_text,modelName):
     if(modelName == "codegen-350M"):
         input_ids = codeGenTokenizer(input_text, return_tensors="pt").input_ids
         generated_ids = codeGenModel.generate(input_ids, max_length=128)
         result = codeGenTokenizer.decode(generated_ids[0], skip_special_tokens=True)
-        return extract_code(result)
+        return extract_code_codegen(result)
     elif(modelName == "mistral-7b"):
         input_ids = mistralTokenizer(generate_prompt_mistral(input_text), return_tensors="pt").input_ids
         generated_ids = mistralModel.generate(input_ids, max_length=128)
         result = mistralTokenizer.decode(generated_ids[0], skip_special_tokens=True)
-        return result
+        return extract_code_mistral(result)
     else:
         return None