timefullytrue/The_Swifty_Chat

1__import__('pysqlite3')2import sys3sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')4 5import os6import gradio as gr7import chromadb8from sentence_transformers import SentenceTransformer9import pandas as pd10import numpy as np11from chromadb.utils import embedding_functions12from openai import OpenAI13 14gem_key = os.getenv("gem_key")15 16dfs = pd.read_csv('TSLyrics.csv')17ids= [str(x) for x in dfs.index.tolist()]18docs = dfs['full_text'].tolist()19client = chromadb.Client()20 21def text_embedding(input):22    model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')23    return model.encode(input)24    25 26collection = client.get_or_create_collection(name="TSlYRICS")27default_ef = text_embedding(docs).tolist()28collection.add(documents=docs,embeddings = default_ef, ids=ids)29 30 31    32def gen_context(query):33    vector = text_embedding(query).tolist()    34    results = collection.query(query_embeddings=vector,n_results=25,include=["documents"])    35    res = "\n".join(str(item) for item in results['documents'][0])36    return res37 38 39 40def get_prompt(history, query):41    prompt = ""42    for i, (old_query, response) in enumerate(history):43        prompt += f"Human: {old_query}\nAI: {response}\n"44    prompt += f"Human: {query}\nAI:"45    return prompt46 47def chat_completion(query):48    49    length = 1000    50    context = gen_context(query)51    52    user_prompt = f"""Based on the context:{context}Answer the below query:{query}"""53    system_prompt = """You are a helpful AI assistant that can answer questions on the context provided in a structured format"""54    final_prompt = f"""<s>[INST]<<SYS>>{system_prompt}<</SYS>>{user_prompt}[/INST]""" 55    56    57    return client.chat.completions.create(model="gemini-2.5-flash",messages=[{"role": "system", "content": system_prompt},{"role": "user","content": user_prompt}]).choices[0].message.content58 59 60client = OpenAI(61    api_key= gem_key,62    base_url="https://generativelanguage.googleapis.com/v1beta/openai/"63)64 65 66demo = gr.Interface(fn=chat_completion, 67                    inputs=[gr.Textbox(label="Query", lines=2)],68                    outputs=[gr.Textbox(label="Result", lines=16)],69                    title="Chat on Taylor Swift Lyrics")70 71demo.queue().launch(share=True)