tiyam-chatbot / app.py
diginoron's picture
Update app.py
ccf44c8 verified
raw
history blame
2.32 kB
import gradio as gr
from sentence_transformers import SentenceTransformer
import pinecone
from transformers import GPT2Tokenizer, GPT2LMHeadModel
import torch
import os
# Load secrets and environment variables
PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY")
PINECONE_INDEX_NAME = os.environ.get("PINECONE_INDEX_NAME")
HF_TOKEN = os.environ.get("HF_TOKEN")
# Step 1: Load embedding model and Pinecone
embedding_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
pinecone.init(api_key=PINECONE_API_KEY)
pc = pinecone.Pinecone(api_key=PINECONE_API_KEY)
index = pc.Index(PINECONE_INDEX_NAME)
# Step 2: Load GPT-2 language model
model_name = "HooshvareLab/gpt2-fa"
tokenizer = GPT2Tokenizer.from_pretrained(model_name, use_auth_token=HF_TOKEN)
model = GPT2LMHeadModel.from_pretrained(model_name, use_auth_token=HF_TOKEN)
model.eval()
# Function: Embed input and search in Pinecone
def retrieve_context(query, top_k=1):
xq = embedding_model.encode(query).tolist()
res = index.query(vector=xq, top_k=top_k, include_metadata=True)
if res.matches:
return res.matches[0].metadata['text']
return ""
# Function: Generate response using GPT-2
def generate_response(query, context):
prompt = f"پرسش: {query}\nپاسخ با توجه به اطلاعات زیر: {context}\nپاسخ:"
input_ids = tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=512)
output_ids = model.generate(input_ids, max_length=256, num_beams=4, no_repeat_ngram_size=2, early_stopping=True)
output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
return output.split("پاسخ:")[-1].strip()
# Gradio interface
def chat(query):
context = retrieve_context(query)
response = generate_response(query, context)
return response
# UI
with gr.Blocks() as demo:
gr.Markdown("## چت‌بات هوشمند تیام\nسوالات خود درباره خدمات دیجیتال مارکتینگ تیام را بپرسید.")
with gr.Row():
inp = gr.Textbox(label="question", placeholder="سوال خود را وارد کنید")
out = gr.Textbox(label="output")
submit = gr.Button("Submit")
submit.click(chat, inputs=inp, outputs=out)
# Launch
if __name__ == "__main__":
demo.launch()