Spaces:
Sleeping
Sleeping
File size: 1,851 Bytes
e75369e 0d7d15e e75369e 328cd2a e75369e 328cd2a e75369e | 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 | import os
from openai import OpenAI
from mistralai import Mistral
from file_utils import load_file
# Embedding API using OpenAI
# Ensure you set your API token in your environment variables or pass it directly
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") # Make sure it's set in your environment
if OPENAI_API_KEY is None:
raise ValueError("Hugging Face API token is missing. Set it as an environment variable: OPENAI_API_KEY")
# Initialize the OpenAI client using the API keys
embedder = OpenAI(api_key=OPENAI_API_KEY)
EMBEDDING_MODEL = "text-embedding-3-large" # "text-embedding-3-small"
def get_embedding(text: str, model: str=EMBEDDING_MODEL) -> list[float]:
"""Get text embeddings from OpenAI."""
result = embedder.embeddings.create(
model=model,
input=text
)
print("Got to get_embedding")
return result.data[0].embedding
# Completions API using HF InferenceClient
# Ensure you set your API token in your environment variables or pass it directly
MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY") # Make sure it's set in your environment
# api_key = os.environ["MISTRAL_API_KEY"]
if MISTRAL_API_KEY is None:
raise ValueError("Hugging Face API token is missing. Set it as an environment variable: MISTRAL_API_KEY")
client = Mistral(api_key=MISTRAL_API_KEY)
COMPLETIONS_MODEL = "mistral-large-latest"
def get_response(messages: list[dict], model: str=COMPLETIONS_MODEL,
temperature=0, max_tokens=800) -> str:
"""Chat completion using Mistral models.
https://docs.mistral.ai/capabilities/completion/
https://docs.mistral.ai/api/#tag/chat
"""
response = client.chat.complete(
model=model,
messages=messages,
max_tokens=max_tokens,
temperature=temperature,
# stream=True
)
return response.choices[0].message.content |