Spaces:
Sleeping
Sleeping
File size: 2,880 Bytes
156fb9c fb13fa4 156fb9c fb13fa4 156fb9c fb13fa4 156fb9c fb13fa4 156fb9c fb13fa4 156fb9c fb13fa4 156fb9c fb13fa4 156fb9c fb13fa4 156fb9c fb13fa4 156fb9c fb13fa4 156fb9c fb13fa4 156fb9c fb13fa4 | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
import torch
# ================= CACHE THE MODEL =================
@st.cache_resource
def load_model():
model_id = "ammoncoder123/IPTchatbotModel1-1.7B" # ← Your correct model repo
# 4-bit quantization for memory efficiency (required for 1.7B on GPU)
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", # Automatically uses GPU if available
torch_dtype=torch.float16,
trust_remote_code=True # Sometimes needed for custom models
)
return pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
max_new_tokens=300,
temperature=0.7,
do_sample=True,
top_p=0.9
)
# Load model once (this will run on first use)
pipe = load_model()
# ==================== CHAT INTERFACE ====================
st.title("IPT Chatbot (1.7B Fine-Tuned Model)")
# Show a disclaimer
st.info("⚠️ This is a small fine-tuned model (1.7B parameters). Answers may contain inaccuracies. Always verify important information.")
# Initialize chat history
if "messages" not in st.session_state:
st.session_state.messages = []
# Display chat history
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 me about IPT, ICT, or anything else..."):
# Add user message
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.markdown(prompt)
# Generate response
with st.chat_message("assistant"):
with st.spinner("Thinking..."):
# Use proper chat format for Instruct models
chat_messages = [
{"role": "user", "content": prompt}
]
outputs = pipe(
chat_messages,
max_new_tokens=300,
temperature=0.7,
do_sample=True,
top_p=0.9
)
# Extract generated text
response = outputs[0]["generated_text"]
# Clean up echoed prompt
if isinstance(response, str) and response.startswith(prompt):
response = response[len(prompt):].strip()
st.markdown(response)
# Save assistant response
st.session_state.messages.append({"role": "assistant", "content": response})
# Optional: Clear chat button
if st.button("Clear Conversation"):
st.session_state.messages = []
st.rerun() |