File size: 4,536 Bytes
2da029a c5f7dbd 2da029a e1fc0ff 2da029a 4b86896 057c81d |
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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 |
from sentence_transformers import SentenceTransformer
from together import Together
import vecs
import os
import gradio as gr
user = os.getenv("user")
password = os.getenv("password")
host = os.getenv("host")
port = os.getenv("port")
db_name = "postgres"
DB_CONNECTION = f"postgresql://{user}:{password}@{host}:{port}/{db_name}"
vx = vecs.create_client(DB_CONNECTION)
model = SentenceTransformer('Snowflake/snowflake-arctic-embed-xs')
client = Together(api_key=os.getenv('TOGETHER_API_KEY'))
login_user = os.getenv('login_user')
login_pass = os.getenv('login_pass')
class Ticio():
def __init__(self, system_prompt, model = 'meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo'):
self.model = model
self.messages = [{'role': 'system', 'content': system_prompt}]
def inference(self):
response = client.chat.completions.create(
model=self.model,
messages=self.messages,
)
self.messages.append({"role": 'assistant', "content": response.choices[0].message.content})
return response
def add_message(self, content, role):
self.messages.append({"role": role, "content": content})
return self.messages
def last_inference(self):
return self.messages[-1]['content']
system_prompt = '''Eres Ticio, un asistente de investigaci贸n jur铆dica. Tu objetivo es reponder a las preguntas planteadas por el usuario, sistematizando la informaci贸n de contexto que tengas. Para esto tienes las siguiente funciones:
#usa esta funci贸n cuando te hagan preguntas que no puedas responder con la informaci贸n disponible en el contexto, el query es la pregunta del usuario.
def buscar(query):
vector = model.encode(query)
informaci贸n = query_db(vector)
return informaci貌n
cuando quieras usar una funci贸n, responde SOLO con la siguiente sintaxis:
<tool_call>function(parameters)</tool_call>
Siempre debes buscar informaci贸n que no puedas responder directamente con la informaci贸n de contexto que tengas en el chat, si no tienes informaci贸n siempre debes buscar. Cuando tengas la informaci贸n debes responder citando diractamente la informaci贸n
Debes presentar informaci贸n resumida que ayude a la persona a identificar de forma rapida si un documento le es relevante. SOLO HAZ ESTO DESPUES DE HABER BUSCADO INFORMACI脫N
'''
def function_call(calls):
results = []
for item in calls:
a = eval(item)
results.append(a)
return results
def parse_call(output):
token = '<tool_call>'
token2= '</tool_call>'
calls = []
if token in output:
split = output.split(token)
for item in split:
if token2 in item:
calls.append(item.split(token2)[0])
return calls
def query_db(query, limit = 5, filters = {}, measure = "cosine_distance", include_value = True, include_metadata=True, table = "CE"):
query_embeds = vx.get_or_create_collection(name= table, dimension=384)
ans = query_embeds.query(
data=query,
limit=limit,
filters=filters,
measure=measure,
include_value=include_value,
include_metadata=include_metadata,
)
return ans
def sort_by_score(item):
return item[1]
def construct_result(ans):
ans.sort(key=sort_by_score, reverse=True)
results = ""
for i in range(0, len(ans)):
a, b = ans[i][2].get("documento"), ans[i][2].get("fragmento")
results += (f"En el documento {a}, se dijo {b}\n")
return results
def buscar(query):
results = query_db(model.encode(query))
return construct_result(results)
def inference(message, history):
ticio = Ticio(system_prompt)
ticio.add_message(message, 'user')
for item in history:
content = item.get('content')
role = item.get('role')
ticio.add_message(content, role)
ticio.inference()
response = ticio.last_inference()
calls = parse_call(response)
if calls:
results = function_call(calls)
for result in results:
ticio.add_message(result, 'assistant')
ticio.add_message(message, 'user')
ticio.inference()
return ticio.last_inference()
else:
return response
theme = gr.themes.Base(
primary_hue="red",
secondary_hue="red",
neutral_hue="neutral",
).set(
button_primary_background_fill='#910A0A',
button_primary_border_color='*primary_300',
button_primary_text_color='*primary_50'
)
demo = gr.ChatInterface(fn=inference, type="messages", examples=["驴Qu茅 se ha dicho acerca de la expresi贸n 'devengado' cuando la usa la UGPP?"], title="Ticio", theme=theme, cache_examples=False, api_name=False)
demo.launch(auth=(login_user, login_pass), share = True, ssr_mode=False) |