nadiamaqbool81 commited on
Commit
6ed2020
·
1 Parent(s): 4679d75

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +75 -0
app.py ADDED
@@ -0,0 +1,75 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from transformers import T5ForConditionalGeneration, AutoTokenizer, RobertaTokenizer,AutoModelForCausalLM,pipeline,TrainingArguments
4
+
5
+
6
+
7
+ models=[
8
+ "nadiamaqbool81/starcoderbase-1b-hf",
9
+ "nadiamaqbool81/starcoderbase-1b-hf_python",
10
+ "nadiamaqbool81/codet5-large-hf",
11
+ "nadiamaqbool81/codet5-large-hf-python",
12
+ "nadiamaqbool81/llama-2-7b-int4-java-code-1.178k",
13
+ "nadiamaqbool81/llama-2-7b-int4-python-code-510"
14
+ ]
15
+ names=[
16
+ "nadiamaqbool81/starcoderbase-java",
17
+ "nadiamaqbool81/starcoderbase-python",
18
+ "nadiamaqbool81/codet5-java",
19
+ "nadiamaqbool81/codet5-python",
20
+ "nadiamaqbool81/llama-2-java",
21
+ "nadiamaqbool81/llama-2-python"
22
+ ]
23
+ model_box=[
24
+ gr.load(f"models/{models[0]}"),
25
+ gr.load(f"models/{models[1]}"),
26
+ gr.load(f"models/{models[2]}"),
27
+ gr.load(f"models/{models[3]}"),
28
+ gr.load(f"models/{models[4]}"),
29
+ gr.load(f"models/{models[5]}"),
30
+ ]
31
+ current_model=model_box[0]
32
+ pythonFlag = "false"
33
+ javaFlag = "false"
34
+
35
+ def the_process(input_text, model_choice):
36
+ global pythonFlag
37
+ global javaFlag
38
+ global output
39
+ if(model_choice == 5):
40
+ if(pythonFlag == "false"):
41
+ tokenizer = AutoTokenizer.from_pretrained("nadiamaqbool81/llama-2-7b-int4-python-code-510")
42
+ model = AutoModelForCausalLM.from_pretrained("nadiamaqbool81/llama-2-7b-int4-python-code-510", load_in_4bit=True, torch_dtype=torch.float16, device_map= {"": 0} )
43
+ output = run_predict(input_text, model, tokenizer)
44
+ pythonFlag = "true"
45
+ elif(pythonFlag == "true"):
46
+ output = run_predict(input_text, model, tokenizer)
47
+ elif(model_choice == 4):
48
+ if(javaFlag == "false"):
49
+ tokenizerJava = AutoTokenizer.from_pretrained("nadiamaqbool81/llama-2-7b-int4-java-code-1.178k")
50
+ modelJava = AutoModelForCausalLM.from_pretrained("nadiamaqbool81/llama-2-7b-int4-java-code-1.178k", load_in_4bit=True, torch_dtype=torch.float16, device_map= {"": 0})
51
+ output = run_predict(input_text, modelJava, tokenizerJava)
52
+ javaFlag = "true"
53
+ elif(javaFlag == "true"):
54
+ output = run_predict(input_text, modelJava, tokenizerJava)
55
+ else:
56
+ a_variable = model_box[model_choice]
57
+ output = a_variable(input_text)
58
+ return(output)
59
+
60
+
61
+ def run_predict(text, model, tokenizer):
62
+ prompt = text
63
+ pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=300)
64
+ result = pipe(f"<s>[INST] {prompt} [/INST]")
65
+ arr = result[0]['generated_text'].split('[/INST]')
66
+ return arr[1]
67
+
68
+
69
+ gr.HTML("""<h1 style="font-weight:600;font-size:50;margin-top:4px;margin-bottom:4px;text-align:center;">Text to Code Generation</h1></div>""")
70
+ model_choice = gr.Dropdown(label="Select Model", choices=[m for m in names], type="index", interactive=True)
71
+ input_text = gr.Textbox(label="Input Prompt")
72
+ output_window = gr.Code(label="Generated Code")
73
+
74
+ interface = gr.Interface(fn=the_process, inputs=[input_text, model_choice], outputs="text")
75
+ interface.launch()