File size: 1,852 Bytes
953bec3 |
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 |
import os
import gradio as gr
import requests
from sentence_transformers import SentenceTransformer
import chromadb
HF_API_URL = os.getenv("HF_API_URL")
HF_TOKEN = os.getenv("HF_TOKEN")
EMB_MODEL = os.getenv("EMB_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
client = chromadb.Client()
collection = client.get_or_create_collection("medini_memory")
embedding_model = SentenceTransformer(EMB_MODEL, use_auth_token=HF_TOKEN)
ROLES = [
"Tutor", "Product Manager", "Project Manager", "Designer/Developer",
"Technical Writer", "Analyst", "Researcher", "Innovator", "Friend", "Video Generator"
]
def retrieve_memory(user_input):
results = collection.query(query_texts=[user_input], n_results=3)
return " ".join(results.get("documents", [[]])[0])
def upsert_memory(user_input, response):
emb = embedding_model.encode(user_input)
collection.add(documents=[user_input], embeddings=[emb], ids=[str(hash(user_input))])
def medini_chat(user_input, role):
context = retrieve_memory(user_input)
payload = {"inputs": f"Role: {role}\nContext: {context}\nUser: {user_input}"}
headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
response = requests.post(HF_API_URL, headers=headers, json=payload)
if response.status_code == 200:
result = response.json()
generated = result.get("generated_text", str(result))
upsert_memory(user_input, generated)
return generated
else:
return f"Error: {response.text}"
iface = gr.Interface(
fn=medini_chat,
inputs=[gr.Textbox(label="Your Message"), gr.Dropdown(ROLES, label="Select Role")],
outputs="text",
title="Medini - Multi-role AI Assistant",
description="Chat with Medini as Tutor, PM, Designer, Analyst, Innovator, Friend, and more."
)
if __name__ == "__main__":
iface.launch(share=True)
|