Create app.py

app.py

@@ -0,0 +1,43 @@
+import torch
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, logging
+
+
+checkpoint = "Salesforce/codegen-350M-mono"
+
+tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(checkpoint, cache_dir="models/", trust_remote_code=True, revision="main")
+
+
+def code_gen(text, max_tokens, temp, top_p, rep_penality):
+    logging.set_verbosity(logging.CRITICAL)
+    pipe = pipeline(
+        model=checkpoint,
+        max_new_tokens=max_tokens,
+        temperature=temp,
+        top_p=top_p,
+        device= "cuda" if torch.cuda.is_available() else "cpu",
+        repetition_penalty=rep_penality
+    )
+
+    response = pipe(text)
+    print(response)
+
+    return response[0]['generated_text']
+
+
+Inferece = gr.Interface(
+    fn=code_gen,
+    inputs=[
+        gr.components.Textbox(label="Input what you want, the AI will make the code for you."),
+        gr.components.Slider(minimum=128, maximum=512, step=128, label="Choose Max Token"),
+        gr.components.Slider(minimum=0.1, maximum=1, step=0.05, label="Choose the model Temperature"),
+        gr.components.Slider(minimum=0.1, maximum=1.25, step=0.05, label="Choose top_p"),
+        gr.components.Slider(minimum=0.1, maximum=2, step=0.1, label="Choose repetition_penalty")
+    ],
+    outputs="text",
+    live=False   # Ensure live is set to False
+)
+
+gr.Markdown('<h2 align="center">AI Code Gen</h2>')
+Inferece.launch()