ITR / app.py
Deepan13's picture
Add CodeLlama demo app
d4417c8
raw
history blame
729 Bytes
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
model_id = "codellama/CodeLlama-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="auto"
)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
def generate_response(prompt):
output = pipe(prompt, max_length=512, do_sample=True, temperature=0.3)[0]['generated_text']
return output
gr.Interface(
fn=generate_response,
inputs=gr.Textbox(lines=5, label="Your prompt"),
outputs=gr.Textbox(label="Code Llama response"),
title="Code Llama Demo",
description="Ask questions or request code snippets!"
).launch()