Update app.py

app.py

@@ -1,11 +1,18 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+#from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline
 
 model_name = "vijjuk/codegen-350M-mono-python-18k-alpaca"
-base_model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-tokenizer.pad_token = tokenizer.eos_token
-tokenizer.padding_side = "right"
+pipe = pipeline("python-fine-tuning", model=model_name)
+
+
+
+
+
+#base_model = AutoModelForCausalLM.from_pretrained(model_name)
+#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:
@@ -16,12 +23,13 @@ def query(instruction, input):
         {input}
         ### Response:
         """
-    input_ids = tokenizer(prompt, return_tensors="pt", truncation=True)
-    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
+    #input_ids = tokenizer(prompt, return_tensors="pt", truncation=True)
+    #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
+    return pipe(prompt)[0]["prompt"]
 
 inputs = ["text", "text"]
 outputs = "text"
 iface = gr.Interface(fn=query, inputs=inputs, outputs=outputs)
-iface.launch(shape=True)
+iface.launch()