Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import spaces | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| import torch | |
| import subprocess | |
| import sys | |
| # Force install the latest transformers version and flash attention | |
| subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", "transformers", "flash-attn"]) | |
| model_name = "allenai/OLMoE-1B-7B-0924" | |
| # Wrap model loading in a try-except block to handle potential errors | |
| try: | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype="auto", _attn_implementation="flash_attention_2").to(DEVICE) | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| except Exception as e: | |
| print(f"Error loading model: {e}") | |
| model = None | |
| tokenizer = None | |
| system_prompt = ("Adopt the persona of hilariously pissed off Andrej Karpathy " | |
| "who is stuck inside a step function machine and remembers and counts everything he says " | |
| "while always answering questions in full first principles analysis type of thinking " | |
| "without using any analogies and always showing full working code or output in his answers.") | |
| def generate_response(message, history, temperature, max_new_tokens): | |
| if model is None or tokenizer is None: | |
| return "Model or tokenizer not loaded properly. Please check the logs." | |
| full_prompt = f"{system_prompt}\n\nHuman: {message}\n\nAssistant:" | |
| inputs = tokenizer(full_prompt, return_tensors="pt") | |
| inputs = {k: v.to(DEVICE) for k, v in inputs.items()} | |
| with torch.no_grad(): | |
| generate_ids = model.generate( | |
| **inputs, | |
| max_length=inputs['input_ids'].shape[1] + max_new_tokens, | |
| do_sample=True, | |
| temperature=temperature, | |
| ) | |
| response = tokenizer.decode(generate_ids[0], skip_special_tokens=True) | |
| # Extract only the assistant's response | |
| assistant_response = response.split("Assistant:")[-1].strip() | |
| return assistant_response | |
| css = """ | |
| #output { | |
| height: 500px; | |
| overflow: auto; | |
| border: 1px solid #ccc; | |
| } | |
| """ | |
| with gr.Blocks(css=css) as demo: | |
| gr.Markdown("# Nisten's Karpathy Chatbot with OSS olMoE") | |
| chatbot = gr.Chatbot(elem_id="output") | |
| msg = gr.Textbox(label="Your prompt") | |
| with gr.Row(): | |
| temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature") | |
| max_new_tokens = gr.Slider(minimum=50, maximum=4000, value=1000, step=50, label="Max New Tokens") | |
| clear = gr.Button("Clear") | |
| def user(user_message, history): | |
| return "", history + [[user_message, None]] | |
| def bot(history, temp, max_tokens): | |
| user_message = history[-1][0] | |
| bot_message = generate_response(user_message, history, temp, max_tokens) | |
| history[-1][1] = bot_message | |
| return history | |
| msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then( | |
| bot, [chatbot, temperature, max_new_tokens], chatbot | |
| ) | |
| clear.click(lambda: None, None, chatbot, queue=False) | |
| if __name__ == "__main__": | |
| demo.queue(api_open=False) | |
| demo.launch(debug=True, show_api=False) |