Update app.py

app.py

@@ -1,25 +1,24 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-# Load your fine-tuned model and tokenizer
-model_name = "EmTpro01/codellama-Code-Generator"  # Use your model name here
+# Load model and tokenizer only once
+model_name = "EmTpro01/codellama-Code-Generator"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
 
-# Function to generate code from a prompt
-def generate_code(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(inputs.input_ids, max_length=150, temperature=0.7, top_k=50)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+# Create pipeline once
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-# Create the Gradio interface
-interface = gr.Interface(
-    fn=generate_code, 
-    inputs="text", 
-    outputs="text", 
-    title="Code Generator",
-    description="Enter a code prompt to generate Python code using the fine-tuned model."
+def generate_response(prompt):
+    # Use the pre-loaded pipeline
+    response = pipe(prompt, max_length=1024, temperature=0.7, top_p=0.95, repetition_penalty=1.15)
+    return response[0]['generated_text']
+
+iface = gr.Interface(
+    fn=generate_response,
+    inputs="text",
+    outputs="text",
+    title="Code Generation Model"
 )
 
-# Launch the app
-interface.launch()
+iface.launch()