File size: 1,545 Bytes
85fbabb 146c71f 85fbabb 146c71f 85fbabb 960e57b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import gradio as gr
model_name = "codellama/CodeLlama-7b-Instruct-hf"
print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
device_map="auto")
print("Model loaded.")
generator = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
temperature=0.1,
top_p=0.95,
max_new_tokens=512,
repetition_penalty=1.05)
def format_prompt(chat):
prompt = ""
for user_msg, ai_reply in chat:
prompt += f"<s>[INST] {user_msg.strip()} [/INST] {ai_reply.strip()}</s>\n"
return prompt
def chat_fn(user_input, history):
history = history or []
prompt = format_prompt(history + [[user_input, ""]])
generated = generator(prompt, do_sample=True)[0]["generated_text"]
answer = generated[len(prompt):].strip()
history.append((user_input, answer))
return "", history
with gr.Blocks() as demo:
gr.Markdown("# 🦙 CodeLlama Copilot\nFree & private code assistant.")
chatbot = gr.Chatbot(label="Developer Assistant", height=400, type="messages")
with gr.Row():
msg = gr.Textbox(placeholder="Ask me coding questions", show_label=False, container=False)
clear = gr.Button("🔄 Clear Conversation")
msg.submit(chat_fn, [msg, chatbot], [msg, chatbot])
clear.click(lambda: ("", []), None, [msg, chatbot])
if __name__ == "__main__":
demo.launch() |