Files
instructor/examples/open_source_examples/runpod.py
T
2023-12-04 23:28:50 -05:00

63 lines
2.1 KiB
Python

import os
import instructor
from openai import OpenAI
from pydantic import BaseModel, Field
from typing import Optional
from instructor import Mode
# Extract API key from environment
runpod_api_key = os.environ.get("RUNPOD_API_KEY")
assert runpod_api_key, "RUNPOD_API_KEY is not set in environment variables"
# Base URL for OpenAI client
runpod_base_url = os.environ.get("RUNPOD_BASE_URL")
assert runpod_base_url, "RUNPOD_BASE_URL is not set in environment variables"
# Initialize OpenAI client
client = instructor.patch(
OpenAI(api_key=runpod_api_key, base_url=runpod_base_url),
mode=Mode.JSON,
)
data = [
"Brandon is 33 years old. He works as a solution architect.",
"Jason is 25 years old. He is the GOAT.",
"Dominic is 45 years old. He is retired.",
"Jenny is 72. She is a wife and a CEO.",
"Holly is 22. She is an explorer.",
"There onces was a prince, named Benny. He ruled for 10 years, which just ended. He started at 22.",
"Simon says, why are you 22 years old marvin?",
]
if __name__ == "__main__":
class UserDetail(BaseModel):
name: str = Field(description="Name extracted from the text")
age: int = Field(description="Age extracted from the text")
occupation: Optional[str] = Field(
default=None, description="Occupation extracted from the text"
)
for content in data:
try:
user = client.chat.completions.create(
response_model=UserDetail,
model="TheBloke_OpenHermes-2.5-Mistral-7B-GPTQ",
messages=[
{
"role": "system",
"content": "You are an expert at outputting json. You output valid JSON.",
},
{
"role": "user",
"content": f"Extract the user details from the following text: {content}. Match your response to the following schema: {UserDetail.model_json_schema()}",
},
],
)
print(f"Result: {user}")
except Exception as e:
print(f"Error: {e}")
continue