Spaces:
Sleeping
Sleeping
File size: 1,871 Bytes
c432fc9 | 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 | import streamlit as st
import json
import numpy as np
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from sentence_transformers import SentenceTransformer
# Load the Mistral model and tokenizer
@st.cache(allow_output_mutation=True)
def load_model():
tokenizer = AutoTokenizer.from_pretrained("mistral7b")
model = AutoModelForSeq2SeqLM.from_pretrained("mistral7b")
return tokenizer, model
# Load Sentence Transformer for embeddings
@st.cache(allow_output_mutation=True)
def load_sentence_transformer():
return SentenceTransformer('all-MiniLM-L6-v2')
tokenizer, model = load_model()
sentence_transformer = load_sentence_transformer()
# Load vector store
@st.cache(allow_output_mutation=True)
def load_vectorstore():
with open('vectorstore.json', 'r') as f:
vectorstore = json.load(f)
return vectorstore
vectorstore = load_vectorstore()
# Function to calculate cosine similarity
def cosine_similarity(vec1, vec2):
return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
# Streamlit UI
st.title("Simple RAG App with Mistral 7B")
query = st.text_input("Enter your question:")
if st.button("Get Answer"):
if query:
# Embed the query
query_embedding = sentence_transformer.encode(query)
# Find the most similar context in the vector store
best_match = max(vectorstore, key=lambda x: cosine_similarity(query_embedding, x['embedding']))
# Generate answer using the Mistral model
inputs = tokenizer.encode(query + " " + best_match['text'], return_tensors='pt')
outputs = model.generate(inputs, max_length=50, num_return_sequences=1)
# Decode and display the answer
answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
st.write("**Answer:**", answer)
else:
st.write("Please enter a question.")
|