07/01/23

app.py

@@ -2,13 +2,7 @@ import re
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model = AutoModelForCausalLM.from_pretrained(
-    'parsak/codegen-350M-mono-lora-instruction',
-)
-tokenizer = AutoTokenizer.from_pretrained('Salesforce/codegen-350M-mono')
 
-tokenizer.pad_token_id = 0  # different to <eos>
-tokenizer.padding_side = "left"  # Allow batched inference
 
 def extract_code(input_text):
     pattern = r"'''py\n(.*?)'''"
@@ -19,28 +13,44 @@ def extract_code(input_text):
     else:
         return None  # Return None if no match is found
 
-def generate_code(input_text):
-    input_ids = tokenizer(input_text, return_tensors="pt").input_ids
-    generated_ids = model.generate(input_ids, max_length=128)
-    result = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
-    return extract_code(result)
+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)
+    elif(modelName == "mistral-7b"):
+        input_ids = mistralTokenizer(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
+    else:
+        return None
+    
 
 def respond(message, chat_history, additional_inputs):
-    return  f"Here's an example code:\n\n```python\n{generate_code(message)}\n```" 
+    return  f"Here's an example code:\n\n```python\n{generate_code(message,additional_inputs)}\n```" 
 
 
+codeGenModel = AutoModelForCausalLM.from_pretrained('parsak/codegen-350M-mono-lora-instruction')
+mistralModel = AutoModelForCausalLM.from_pretrained('parsak/mistral-code-7b-instruct')
+codeGenTokenizer = AutoTokenizer.from_pretrained('Salesforce/codegen-350M-mono')
+mistralTokenizer = AutoTokenizer.from_pretrained('parsak/mistral-code-7b-instruct')
+codeGenTokenizer.pad_token_id = 0
+codeGenTokenizer.padding_side = "left" 
 
 
+dropdown = gr.Dropdown(label="Models",choices=["codegen-350M", "mistral-7b"], value="codegen-350M")
 
-
-gr.ChatInterface(respond,
+interface = gr.ChatInterface(respond,
         retry_btn= gr.Button(value="Retry"), 
         undo_btn=None, clear_btn=gr.Button(value="Clear"),
         additional_inputs=[
-            gr.Dropdown(["annen", "baban"])
+            dropdown
         ]
-        ).launch()
-
+        )
 
 
+if __name__ == "__main__":
+    interface.launch()