| 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() |