Spaces:
Sleeping
Sleeping
File size: 2,123 Bytes
6f8b2d3 |
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 |
from typing import List
from langchain_ollama import OllamaLLM
from langchain.output_parsers import PydanticOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field
llm = OllamaLLM(model="gemma2", temperature=0.8)
class Person(BaseModel):
name: str = Field(description="Person's name")
age: int = Field(description="Person's age")
nationality: str = Field(description="Person's nationality")
gender: str = Field(description="Person's gender")
description: str = Field(description="Person's description")
occupation: str = Field(description="Person's occupation")
known_languages: str = Field(
description="Person's known languages it could be one or more")
class Witness(Person):
testimony: str = Field(description="Witness's testimony")
class Suspect(Person):
guilty: bool = Field(description="Suspect's guilt")
background: str = Field(
description="Suspect's background during the crime")
class Crime(BaseModel):
suspects: List[Suspect]
witnesses: List[Witness]
event: str = Field(
description="A long description with all the details about the crime with date, aproximetly time and place")
date: str = Field(description="Crime date")
location: str = Field(description="Crime location")
summary: str = Field(description="Brief description of the crime")
motive: str = Field(description="Motive of the crime")
parser = PydanticOutputParser(pydantic_object=Crime)
prompt = PromptTemplate(
template="Create a crime in {country} in all the text should be in the language of this country, and a list with {suspects} suspects and {witnesses} where only one of the suspects is guilty. \n {format_instructions}",
input_variables=["suspects", "witnesses"],
partial_variables={
"format_instructions": parser.get_format_instructions()},
)
def generate_crime(suspects, witnesses, country):
chain = prompt | llm | parser
output = chain.invoke(
{"suspects": suspects, "witnesses": witnesses, "country": country})
return output
|