|
|
import streamlit as st |
|
|
import chromadb |
|
|
import requests |
|
|
import os |
|
|
|
|
|
|
|
|
MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.1" |
|
|
API_URL = f"https://api-inference.huggingface.co/models/{MODEL_ID}" |
|
|
|
|
|
API_TOKEN = st.secrets["HUGGINGFACEHUB_API_TOKEN"] |
|
|
|
|
|
|
|
|
headers = { |
|
|
"Authorization": f"Bearer {API_TOKEN}" |
|
|
} |
|
|
|
|
|
|
|
|
chroma_client = chromadb.PersistentClient(path="chroma_store") |
|
|
collection = chroma_client.get_or_create_collection(name="tech_docs") |
|
|
|
|
|
|
|
|
def query_huggingface(prompt): |
|
|
payload = { |
|
|
"inputs": prompt, |
|
|
"options": {"wait_for_model": True} |
|
|
} |
|
|
response = requests.post(API_URL, headers=headers, json=payload) |
|
|
|
|
|
|
|
|
if response.status_code != 200: |
|
|
return f"❌ HF API Error: {response.status_code} - {response.text}" |
|
|
|
|
|
try: |
|
|
result = response.json() |
|
|
if isinstance(result, list) and "generated_text" in result[0]: |
|
|
return result[0]["generated_text"] |
|
|
elif isinstance(result, dict) and "error" in result: |
|
|
return f"❌ Model Error: {result['error']}" |
|
|
else: |
|
|
return "⚠️ Unexpected model response format." |
|
|
except Exception as e: |
|
|
return f"⚠️ Failed to parse response: {e}" |
|
|
|
|
|
|
|
|
st.title("💬 Ask Me Anything - Tech RAG Chatbot") |
|
|
|
|
|
user_query = st.text_input("🔎 Ask your question:") |
|
|
|
|
|
if user_query: |
|
|
|
|
|
results = collection.query(query_texts=[user_query], n_results=3) |
|
|
context = "\n".join(results["documents"][0]) if results["documents"] else "" |
|
|
|
|
|
|
|
|
prompt = f"""Answer the question using the context below: |
|
|
|
|
|
Context: |
|
|
{context} |
|
|
|
|
|
Question: |
|
|
{user_query} |
|
|
|
|
|
Answer:""" |
|
|
|
|
|
|
|
|
with st.spinner("Thinking..."): |
|
|
answer = query_huggingface(prompt) |
|
|
|
|
|
st.markdown("### 📢 Answer:") |
|
|
st.write(answer) |
|
|
|