Spaces:
Sleeping
Sleeping
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,101 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import streamlit as st
|
2 |
+
import os
|
3 |
+
import numpy as np
|
4 |
+
from sentence_transformers import SentenceTransformer
|
5 |
+
import faiss
|
6 |
+
from openai import OpenAI
|
7 |
+
from PIL import Image
|
8 |
+
|
9 |
+
class IntegratedChatSystem:
|
10 |
+
def __init__(self, api_key: str, model_name: str, embedding_dim: int = 384):
|
11 |
+
self.api_key = api_key
|
12 |
+
self.model_name = model_name
|
13 |
+
self.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
|
14 |
+
self.embedding_dim = embedding_dim
|
15 |
+
self.index = faiss.IndexFlatIP(embedding_dim)
|
16 |
+
self.metadata = []
|
17 |
+
self.client = OpenAI(api_key=api_key)
|
18 |
+
|
19 |
+
def _add_to_index(self, vector: np.ndarray, metadata: dict):
|
20 |
+
self.index.add(vector)
|
21 |
+
self.metadata.append(metadata)
|
22 |
+
|
23 |
+
def add_image(self, image_path: str, context_text: str):
|
24 |
+
filename = os.path.basename(image_path)
|
25 |
+
if not os.path.exists(image_path):
|
26 |
+
raise FileNotFoundError(f"Image not found: {image_path}")
|
27 |
+
embedding = self.embedding_model.encode(context_text)
|
28 |
+
embedding = np.expand_dims(embedding, axis=0)
|
29 |
+
self._add_to_index(embedding, {"filepath": filename, "context": context_text})
|
30 |
+
|
31 |
+
def chat(self, user_message: str, similarity_threshold: float = 0.7, top_k: int = 3):
|
32 |
+
message_embedding = self.embedding_model.encode(user_message)
|
33 |
+
message_embedding = np.expand_dims(message_embedding, axis=0)
|
34 |
+
distances, indices = self.index.search(message_embedding, top_k)
|
35 |
+
relevant_images = [
|
36 |
+
self.metadata[i] for i, distance in zip(indices[0], distances[0])
|
37 |
+
if i != -1 and distance >= similarity_threshold
|
38 |
+
]
|
39 |
+
|
40 |
+
system_prompt = """You are an assistant chatbot. You should help the user by answering their question."""
|
41 |
+
|
42 |
+
enhanced_message = user_message
|
43 |
+
if relevant_images:
|
44 |
+
image_contexts = "\n".join(f"- {img['context']}" for img in relevant_images)
|
45 |
+
enhanced_message = f"{user_message}\n\nContext from relevant images:\n{image_contexts}"
|
46 |
+
|
47 |
+
try:
|
48 |
+
completion = self.client.chat_completions.create(
|
49 |
+
model=self.model_name,
|
50 |
+
messages=[
|
51 |
+
{"role": "system", "content": system_prompt},
|
52 |
+
{"role": "user", "content": enhanced_message}
|
53 |
+
]
|
54 |
+
)
|
55 |
+
|
56 |
+
response = completion.choices[0].message.content
|
57 |
+
|
58 |
+
return {
|
59 |
+
"response": response,
|
60 |
+
"images": relevant_images if relevant_images else None
|
61 |
+
}
|
62 |
+
|
63 |
+
except Exception as e:
|
64 |
+
print(f"Error calling OpenAI API: {str(e)}")
|
65 |
+
return {
|
66 |
+
"response": "I apologize, but I encountered an error processing your request.",
|
67 |
+
"images": None
|
68 |
+
}
|
69 |
+
|
70 |
+
# Initialize the chat system
|
71 |
+
api_key = "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"
|
72 |
+
model_name = "ft:gpt-3.5-turbo-0125:brenin::AlVMkeUb"
|
73 |
+
chat_system = IntegratedChatSystem(api_key, model_name)
|
74 |
+
|
75 |
+
# Add images
|
76 |
+
image_folder = "images"
|
77 |
+
chat_system.add_image(os.path.join(image_folder, "sequence diagram.png"), "A diagram showing the sequence of how it is supposed to work")
|
78 |
+
chat_system.add_image(os.path.join(image_folder, "UX workflow.png"), "A flowchart of showing the UX workflow")
|
79 |
+
|
80 |
+
# Streamlit UI
|
81 |
+
st.title("Chat with Integrated Image Context")
|
82 |
+
|
83 |
+
st.sidebar.title("Chat System")
|
84 |
+
user_message = st.text_input("Your message:", placeholder="Type your message here...")
|
85 |
+
|
86 |
+
if st.button("Send"):
|
87 |
+
if user_message.strip():
|
88 |
+
result = chat_system.chat(user_message)
|
89 |
+
st.write(f"**Assistant:** {result['response']}")
|
90 |
+
|
91 |
+
if result["images"]:
|
92 |
+
st.write("Relevant Images:")
|
93 |
+
for img in result["images"]:
|
94 |
+
image_path = os.path.join(image_folder, img["filepath"])
|
95 |
+
if os.path.exists(image_path):
|
96 |
+
st.image(Image.open(image_path), caption=img["context"])
|
97 |
+
else:
|
98 |
+
st.write(f"Image not found: {img['filepath']}")
|
99 |
+
else:
|
100 |
+
st.error("Please enter a message.")
|
101 |
+
|