Spaces:
Sleeping
Sleeping
| import os | |
| import torch | |
| import gradio as gr | |
| from model import SmolLM | |
| from huggingface_hub import hf_hub_download | |
| device = "mps" if torch.backends.mps.is_available() else "cpu" | |
| hf_token = os.environ.get("HF_TOKEN") | |
| repo_id = "ZivK/smollm2-end-of-sentence" | |
| model_options = { | |
| "Word-level Model": "word_model.ckpt", | |
| "Token-level Model": "token_model.ckpt" | |
| } | |
| label_map = {0: "Incomplete", 1: "Complete"} | |
| models = {} | |
| for model_name, filename in model_options.items(): | |
| print(f"Loading {model_name} ...") | |
| checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename, token=hf_token) | |
| models[model_name] = SmolLM.load_from_checkpoint(checkpoint_path).to(device) | |
| models[model_name].eval() | |
| def classify_sentence(sentence, model_choice): | |
| model = models[model_choice] | |
| inputs = model.tokenizer(sentence, return_tensors="pt", padding=True, truncation=True).to(device) | |
| with torch.no_grad(): | |
| logits = model(inputs) | |
| confidence = torch.sigmoid(logits).item() * 100 | |
| predicted_class = 1 if confidence > 50.0 else 0 | |
| return label_map[predicted_class], confidence | |
| def chatbot_reply(history, user_input, model_choice): | |
| classification, confidence = classify_sentence(user_input, model_choice) | |
| if classification == "Incomplete": | |
| bot_message = "It looks like you may have stopped mid-sentence. Please finish your thought! Confidence: " + \ | |
| f"{(100.0-confidence):.2f}" | |
| else: | |
| bot_message = f"Thank you for sharing a complete sentence! Confidence: {confidence:.2f}" | |
| # Append the user message and bot response to the conversation history | |
| history.append((user_input, bot_message)) | |
| return history, "" | |
| with gr.Blocks() as demo: | |
| gr.Markdown( | |
| "## Sentence Completeness Chatbot\nType a message and see if the model thinks it’s complete or incomplete!") | |
| gr.Markdown("#### [Click here to view the model on Hugging Face](https://huggingface.co/ZivK/smollm2-end-of-sentence)") | |
| # 3. Create a stateful Chatbot plus an input textbox | |
| chatbot = gr.Chatbot(label="Chat with Me!") | |
| state = gr.State([]) # This will store the conversation history | |
| with gr.Row(): | |
| user_input = gr.Textbox(show_label=False, placeholder="Type your sentence here...") | |
| submit_btn = gr.Button("Submit") | |
| with gr.Row(): | |
| model_input = gr.Dropdown(choices=list(model_options.keys()), label="Select Model") | |
| # 4. Bind the chatbot function | |
| submit_btn.click(fn=chatbot_reply, | |
| inputs=[state, user_input, model_input], | |
| outputs=[chatbot, user_input]) | |
| # We also want pressing Enter to do the same as clicking submit | |
| user_input.submit(fn=chatbot_reply, | |
| inputs=[state, user_input, model_input], | |
| outputs=[chatbot, user_input]) | |
| # Launch the demo | |
| demo.launch() | |