Update app.py

Updated the logic

app.py

@@ -1,7 +1,31 @@
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+model_name = "Salesforce/codegen-350M-mono"
+base_model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config, use_cache = False, device_map=device_map)
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = "right"
+
+
+
+
+def query(instruction, input):
+    prompt = f"""### Instruction:
+        Use the Task below and the Input given to write the Response, which is a programming code that can solve the Task.
+        ### Task:
+        {instruction}
+        ### Input:
+        {input}
+        ### Response:
+        """
+    input_ids = tokenizer(prompt, return_tensors="pt", truncation=True).input_ids.cuda()
+    output_base = base_model.generate(input_ids=input_ids, max_new_tokens=500, do_sample=True, top_p=0.9,temperature=0.5)
+    response = "{tokenizer.batch_decode(output_base.detach().cpu().numpy(), skip_special_tokens=True)[0][len(prompt):]}"
+    return response
+
+inputs = ["text", "text"]
+outputs = "text"
+iface = gr.Interface(fn=query, inputs="text", outputs="text")
 iface.launch()