Update app.py

app.py

@@ -2,15 +2,15 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 
-# Load the fine-tuned MFTCoder model
+# Load the CodeLlama model and tokenizer
 @st.cache_resource()
 def load_model():
-    MODEL_NAME = "path-to-your-finetuned-model"  # Replace with your MFTCoder fine-tuned model path
+    MODEL_NAME = "codellama/CodeLlama-7b-hf"  # Model name
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        torch_dtype=torch.float16,  # Use float16 for performance optimization
-        device_map="auto"  # Automatically allocate to CPU/GPU
+        MODEL_NAME, 
+        torch_dtype=torch.float32,  # Use float32 for CPU
+        device_map="cpu"            # Map model to CPU
     )
     return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
@@ -18,31 +18,31 @@ def load_model():
 code_generator = load_model()
 
 # Streamlit UI
-st.title("MFTCoder-powered Code Bot 🚀")
-st.subheader("Generate high-quality code snippets with fine-tuned CodeLlama!")
+st.title("CodeLlama-7B Code Bot 🚀")
+st.subheader("Generate code snippets using CodeLlama-7b-hf on CPU")
 
 # User input
-prompt = st.text_area("Enter a code prompt to generate code:")
+prompt = st.text_area("Enter a coding prompt (e.g., 'Write a Python function to sort a list'): ")
 
-# Generate code
+# Generate Code
 if st.button("Generate Code"):
     if prompt.strip():
         st.info("Generating code... Please wait ⏳")
         try:
-            # Generate code using the fine-tuned MFTCoder model
+            # Generate code using the CodeLlama model
             response = code_generator(
                 prompt,
-                max_length=256,  # Adjust as needed
-                temperature=0.3,  # Lower temperature for accurate outputs
-                num_return_sequences=1,
-                do_sample=True
+                max_length=512,  # Increase for longer code generation
+                temperature=0.2,  # Lower temperature for more deterministic results
+                do_sample=True,   # Enable sampling
+                num_return_sequences=1
             )
             generated_code = response[0]['generated_text']
             # Display the code output
-            st.code(generated_code, language="python")  # Default to Python for generated output
+            st.code(generated_code, language="python")  # Change language as needed
         except Exception as e:
             st.error(f"Error: {str(e)}")
     else:
         st.warning("Please enter a prompt.")
 
-st.caption("Created by Shamil")
+st.caption("Powered by CodeLlama-7B | Streamlit UI | CPU Optimized")