Spaces:
Sleeping
Sleeping
File size: 4,936 Bytes
f3dd217 |
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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
import streamlit as st
from transformers import pipeline
import torch
import time
from typing import List, Dict
import functools
import signal
class TimeoutError(Exception):
pass
def timeout(seconds):
def decorator(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
def handler(signum, frame):
raise TimeoutError(f"Function call timed out after {seconds} seconds")
# Set the timeout handler
signal.signal(signal.SIGALRM, handler)
signal.alarm(seconds)
try:
result = func(*args, **kwargs)
finally:
# Disable the alarm
signal.alarm(0)
return result
return wrapper
return decorator
class SourceVerifier:
def __init__(self):
self.sources: List[Dict] = []
def add_source(self, text: str, metadata: Dict) -> None:
self.sources.append({"content": text, "metadata": metadata})
def verify_statement(self, statement: str) -> Dict:
matches = []
for source in self.sources:
if any(word.lower() in source["content"].lower()
for word in statement.split()):
matches.append(source)
return {
"verified": len(matches) > 0,
"matches": matches,
"confidence": len(matches) / len(self.sources) if self.sources else 0
}
@st.cache_resource(show_spinner=False)
def load_pipeline():
try:
return pipeline(
"text-generation",
model="sshleifer/tiny-gpt2", # Tiny 2M parameter model
device="cpu", # Force CPU usage
model_kwargs={"low_memory": True}
)
except Exception as e:
st.error(f"Failed to load model: {str(e)}")
return None
@timeout(10) # 10 second timeout
def generate_response(generator, prompt: str) -> str:
try:
result = generator(
prompt,
max_length=50, # Short response
num_return_sequences=1,
temperature=0.7,
do_sample=True,
)
return result[0]['generated_text']
except TimeoutError:
return "Response generation timed out. Please try again."
except Exception as e:
return f"Error generating response: {str(e)}"
def init_page():
st.set_page_config(
page_title="Quick Chat Demo",
page_icon="💬",
layout="centered"
)
st.title("Quick Chat Demo")
if "messages" not in st.session_state:
st.session_state.messages = [
{"role": "assistant", "content": "Hi! I'm a simple chat demo. How can I help?"}
]
if "verifier" not in st.session_state:
st.session_state.verifier = SourceVerifier()
def handle_file_upload():
uploaded_file = st.file_uploader("Upload source document", type=["txt", "md", "json"])
if uploaded_file:
try:
content = uploaded_file.read().decode()
st.session_state.verifier.add_source(
content,
{"filename": uploaded_file.name, "type": uploaded_file.type}
)
st.success(f"Added source: {uploaded_file.name}")
except Exception as e:
st.error(f"Error processing file: {str(e)}")
def main():
init_page()
# Load the model with a progress bar
with st.spinner("Loading (should take < 5 seconds)..."):
generator = load_pipeline()
if generator is None:
st.error("Failed to initialize chat. Please refresh the page.")
return
# Sidebar for document upload
with st.sidebar:
st.header("Sources")
handle_file_upload()
# Display existing 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("Say something"):
# Add user message
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.write(prompt)
# Generate response with timeout
with st.chat_message("assistant"):
with st.spinner("Responding..."):
response = generate_response(generator, prompt)
verification = st.session_state.verifier.verify_statement(response)
st.write(response)
if verification["verified"]:
with st.expander("Sources"):
st.json(verification)
st.session_state.messages.append({
"role": "assistant",
"content": response,
"verification": verification
})
if __name__ == "__main__":
main() |