Spaces:
Sleeping
Sleeping
File size: 1,507 Bytes
428116e c89e432 428116e c89e432 428116e c89e432 428116e |
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 |
import gradio as gr
import sqlite_utils
import llm
import pandas as pd
from openai import OpenAI
import os
if not os.getenv("OPENAI_API_KEY"):
raise ValueError("OPENAI_API_KEY must be set")
def retrieval(question: str):
db = sqlite_utils.Database("metropolen.db")
embedding_model = llm.get_embedding_model("ada-002")
collection = llm.Collection("sections", db, model=embedding_model)
data = [(entry.id, entry.content, entry.score) for entry in collection.similar(question, number=3)]
df = pd.DataFrame(data, columns=['section_heading', 'content', 'score'])
documents = (df.content).str.cat(sep='\n')
return documents
def make_prompt(documents, question):
prompt = f"""Given the context below, answer the question at the end of the text.
---
{documents}
---
{question}
"""
return prompt
def generation(prompt, api_key=None, model="gpt-4o"):
api_key = api_key if api_key else os.getenv('OPENAI_API_KEY')
client = OpenAI(api_key=api_key)
c = client.chat.completions.create(model=model,
messages=[{"role": "user", "content": prompt}])
return c.choices[0].message.content
def rag(frage: str) -> str:
documents = retrieval(question=frage)
prompt = make_prompt(documents=documents, question=frage)
antwort = generation(prompt=prompt, api_key=os.getenv('OPENAI_API_KEY'), model="gpt-4o")
return antwort
demo = gr.Interface(fn=rag, inputs="text", outputs="text")
demo.launch()
|