Spaces:
Sleeping
Sleeping
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() | |