rohangbs commited on
Commit
fc6ab5e
·
verified ·
1 Parent(s): 8c15bc1

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +101 -0
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
+