Spaces:
Sleeping
Sleeping
File size: 3,163 Bytes
c755398 |
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 92 93 94 95 96 97 98 99 100 101 |
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Set page config
st.set_page_config(
page_title="Phi2 GPro Chat",
page_icon="🚀",
layout="wide"
)
# Initialize session state for chat history if it doesn't exist
if "messages" not in st.session_state:
st.session_state.messages = []
@st.cache_resource
def load_model():
from peft import PeftModel, PeftConfig
# Load base model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
base_model = AutoModelForCausalLM.from_pretrained(
"microsoft/phi-2",
torch_dtype=torch.float16,
device_map="auto",
low_cpu_mem_usage=True
)
# Load and apply adapter weights
model = PeftModel.from_pretrained(base_model, "sft-model")
return model, tokenizer
def generate_response(prompt, model, tokenizer, max_length=512):
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=max_length)
inputs = {k: v.to(model.device) for k, v in inputs.items()}
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=512,
do_sample=True,
temperature=0.7,
top_p=0.9,
pad_token_id=tokenizer.eos_token_id
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Remove the input prompt from the response
response = response[len(tokenizer.decode(inputs['input_ids'][0], skip_special_tokens=True)):].strip()
return response
# Main app
st.title("Phi2 GPro Chat 🚀")
# Load model
try:
model, tokenizer = load_model()
st.success("Model loaded successfully! Ready to chat.")
except Exception as e:
st.error(f"Error loading model: {str(e)}")
st.stop()
# Display chat messages
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.write(message["content"])
# Chat input
if prompt := st.chat_input("What would you like to discuss?"):
# Add user message to chat history
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.write(prompt)
# Generate response
with st.chat_message("assistant"):
with st.spinner("Thinking..."):
# Prepare conversation history
conversation = ""
for msg in st.session_state.messages:
if msg["role"] == "user":
conversation += f"User: {msg['content']}\n"
else:
conversation += f"Assistant: {msg['content']}\n"
response = generate_response(conversation, model, tokenizer)
st.write(response)
st.session_state.messages.append({"role": "assistant", "content": response})
# Add a sidebar with information
with st.sidebar:
st.title("About")
st.markdown("""
This is a chatbot powered by the Phi2 GPro model.
Feel free to ask questions and engage in conversation!
**Features:**
- Contextual responses
- Memory of conversation
- Fast inference
""") |