CS269-codellama / app.py
Vaishnavi-15's picture
Update app.py
9d2f6bd verified
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
#from huggingface_hub import login
#login(token="hf_VExbFezQQyzOnbpBoRgNxXjiRfMFTGUyj")
my_token="hf_VExbFezQQyzOnbpBoRgNxXjiRfMFTGUyj"
# Load model and tokenizer
model_name = "meta-llama/CodeLlama-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name, token=my_token)
model = AutoModelForCausalLM.from_pretrained(model_name, token=my_token)
# Define the inference function
def generate_code(prompt):
# Tokenize the input prompt
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)
# Generate code using the model
outputs = model.generate(inputs["input_ids"], max_length=100, num_return_sequences=1)
# Decode the generated output to a string
generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
return generated_code
# Create Gradio interface
interface = gr.Interface(
fn=generate_code,
inputs="text",
outputs="text",
title="CodeLlama-7b Python Code Generator",
description="Generate Python code using the CodeLlama-7b model. Simply input a prompt and get back the generated code.",
)
# Launch the Gradio interface
interface.launch()