Update app.py

app.py

@@ -3,27 +3,23 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
 
 # Check if the token is being accessed
-hf_token = os.getenv("HF_HOME")
-if hf_token:
-    print("Successfully retrieved Hugging Face token.")
-else:
-    print("Failed to retrieve Hugging Face token.")
+hf_token = os.environ.get("HF_HOME", None)
 
-# # Load the model and tokenizer
-# model_name = "meta-llama/CodeLlama-7b-hf"
-# model = AutoModelForCausalLM.from_pretrained(model_name)
-# tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Load the model and tokenizer
+model_name = "meta-llama/CodeLlama-7b-hf"
+model = AutoModelForCausalLM.from_pretrained(model_name, token=hf_token)
+tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
 
 
-# def generate_code(prompt):
-#     inputs = tokenizer(prompt, return_tensors="pt")
-#     outputs = model.generate(inputs["input_ids"], max_length=200)
-#     code = tokenizer.decode(outputs[0], skip_special_tokens=True)
-#     return code
+def generate_code(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(inputs["input_ids"], max_length=200)
+    code = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return code
 
-# # Set up the Gradio interface
-# demo = gr.Interface(fn=generate_code, 
-#              inputs="text", 
-#              outputs="text", 
-#              title="CodeLlama 7B Model",
-#              description="Generate code with CodeLlama-7b-hf.").launch()
+# Set up the Gradio interface
+demo = gr.Interface(fn=generate_code, 
+             inputs="text", 
+             outputs="text", 
+             title="CodeLlama 7B Model",
+             description="Generate code with CodeLlama-7b-hf.").launch()