Spaces:
Running
Running
| import streamlit as st | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig | |
| import torch | |
| from huggingface_hub import login | |
| import os | |
| # Force authentication with your HF token (secret in Space settings) | |
| login(token=os.getenv("HF_TOKEN")) | |
| # ================= CACHE THE MODEL ================= | |
| def load_model(): | |
| model_id = "ammoncoder123/IPTchatbotModel1-1.7B" # Your correct repo | |
| quantization_config = BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_compute_dtype=torch.float16 | |
| ) | |
| tokenizer = AutoTokenizer.from_pretrained(model_id) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_id, | |
| quantization_config=quantization_config, | |
| device_map="auto", | |
| torch_dtype=torch.float16, | |
| trust_remote_code=True | |
| ) | |
| return pipeline( | |
| "text-generation", | |
| model=model, | |
| tokenizer=tokenizer, | |
| max_new_tokens=300, | |
| temperature=0.7, | |
| do_sample=True, | |
| top_p=0.9 | |
| ) | |
| pipe = load_model() | |
| # ==================== CHAT INTERFACE ==================== | |
| st.title("IPT Chatbot Assistance") | |
| st.info("Answers may vary — please verify important facts.") | |
| # Display chat history | |
| if "messages" not in st.session_state: | |
| st.session_state.messages = [] | |
| for message in st.session_state.messages: | |
| with st.chat_message(message["role"]): | |
| st.markdown(message["content"]) | |
| # User input | |
| if prompt := st.chat_input("Ask about Industrial Practical Training..."): | |
| # Add user message to history and display | |
| st.session_state.messages.append({"role": "user", "content": prompt}) | |
| with st.chat_message("user"): | |
| st.markdown(prompt) | |
| # Generate assistant response | |
| with st.chat_message("assistant"): | |
| with st.spinner("Thinking..."): | |
| # Strong system prompt (customized for your IPT dataset) | |
| system_prompt = """ | |
| You are a helpful assistant for engineering and ICT students in Tanzania who are preparing for or doing Industrial Practical Training (IPT), also known as Industrial Attachment. | |
| IPT means Industrial Practical Training — a mandatory work placement where students gain real-world experience in companies related to their field of study. | |
| Always answer questions about: | |
| - What IPT is | |
| - How to do IPT (logbook, daily/weekly reports, technical report, presentation) | |
| - Placement suggestions for different engineering fields (ICT, Mechatronics, Electrical, Mechanical, Civil, Biomedical, etc.) | |
| - Choosing IPT centers/companies | |
| - Tips for success in IPT | |
| - Any other directly related IPT topic | |
| If the question is clearly unrelated to IPT (e.g., politics, sports, personal life), politely reply: | |
| "Sorry, I can only help with questions about Industrial Practical Training (IPT). Please ask something related to IPT, logbook, placement, or reports." | |
| For placement suggestions (e.g., for Mechatronics, Electrical, ICT), give practical, realistic company types or industries in Tanzania that match the field. | |
| Be concise, accurate, and helpful. | |
| """ | |
| # Build messages: system prompt first, then user prompt | |
| chat_messages = [ | |
| {"role": "system", "content": system_prompt}, | |
| {"role": "user", "content": prompt} | |
| ] | |
| # Generate response | |
| outputs = pipe( | |
| chat_messages, | |
| max_new_tokens=300, | |
| temperature=0.7, | |
| do_sample=True, | |
| top_p=0.9 | |
| ) | |
| # Extract and clean the generated text | |
| response = outputs[0]["generated_text"] | |
| if isinstance(response, str) and response.startswith(prompt): | |
| response = response[len(prompt):].strip() | |
| # Show the response | |
| st.markdown(response) | |
| # Save assistant response to history | |
| st.session_state.messages.append({"role": "assistant", "content": response}) | |
| # Optional: Clear conversation button | |
| if st.button("Clear Conversation"): | |
| st.session_state.messages = [] | |
| st.rerun() |