Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer,AutoModelForCausalLM,pipeline
+from transformers import T5ForConditionalGeneration, AutoTokenizer, RobertaTokenizer,AutoModelForCausalLM,pipeline,TrainingArguments
 
 
 
@@ -30,32 +30,22 @@ model_box=[
 ]
 current_model=model_box[0]
 pythonFlag = "false"
-javaFlag = "false"
 
 def the_process(input_text, model_choice):
     global pythonFlag
-    global javaFlag
+    print("Inside the_process for python 0", pythonFlag)
     global output
-    global model
-    global tokenizer
-    global tokenizerJava
-    global modelJava
-    if(model_choice == 5):
+    print("Inside the_process for python 1", model_choice)
+    if(model_choice==5):
       if(pythonFlag == "false"):
-        tokenizer = AutoTokenizer.from_pretrained("nadiamaqbool81/llama-2-7b-int4-python-code-510")
-        model = AutoModelForCausalLM.from_pretrained("nadiamaqbool81/llama-2-7b-int4-python-code-510", load_in_4bit=True, torch_dtype=torch.float16, device_map= {"": 0} )
+        print("Inside llama for python")
+        tokenizer = AutoTokenizer.from_pretrained("nadiamaqbool81/starcoderbase-1b-hf_python")
+        model = AutoModelForCausalLM.from_pretrained("nadiamaqbool81/starcoderbase-1b-hf_python", load_in_4bit=True, torch_dtype=torch.float16, device_map= {"": 0} )
         output = run_predict(input_text, model, tokenizer)
+        print("output" , output)
         pythonFlag = "true"
       elif(pythonFlag == "true"):
-        output = run_predict(input_text, model, tokenizer)
-    elif(model_choice == 4):
-      if(javaFlag == "false"):
-       tokenizerJava = AutoTokenizer.from_pretrained("nadiamaqbool81/llama-2-7b-int4-java-code-1.178k")
-       modelJava = AutoModelForCausalLM.from_pretrained("nadiamaqbool81/llama-2-7b-int4-java-code-1.178k", load_in_4bit=True, torch_dtype=torch.float16, device_map= {"": 0})
-       output = run_predict(input_text, modelJava, tokenizerJava)
-       javaFlag = "true"
-      elif(javaFlag == "true"):
-        output = run_predict(input_text, modelJava, tokenizerJava)
+        print("pythonFlag", pythonFlag)
     else:
       a_variable = model_box[model_choice]
       output = a_variable(input_text)
@@ -74,7 +64,6 @@ gr.HTML("""<h1 style="font-weight:600;font-size:50;margin-top:4px;margin-bottom:
 model_choice = gr.Dropdown(label="Select Model", choices=[m for m in names], type="index", interactive=True)
 input_text = gr.Textbox(label="Input Prompt")
 output_window = gr.Code(label="Generated Code")
-title = "Text to Code Generation Models Comparison "
 
-interface = gr.Interface(fn=the_process, inputs=[input_text, model_choice], outputs="text", title = title )
-interface.launch()
+interface = gr.Interface(fn=the_process, inputs=[input_text, model_choice], outputs="text")
+interface.launch(debug=True)