spec_generation.py

from pydantic import BaseModel, Field
import image_generation, json
from vllm import SamplingParams
from vllm.sampling_params import StructuredOutputsParams

class SpecificationsOutput(BaseModel):
    reasoning: str = Field(description="A brief analysis (2-3 sentences) justifying the chosen chart and data")
    subject: str = Field(description="The subject that is addressed by the teacher")
    chart_type: str = Field(description="The chosen chart type")
    data: str = Field(description="The chosen data including numeric values, labels")
    task: str = Field(description="A single self-contained description combining the chart type and the exact data values and labels, detailed enough to generate the plot with no other context (e.g. 'A bar chart of monthly sales: Jan=120, Feb=90, ...')")
    
SPEC_PROMPT = """
You have the role to design an educative figure so that students can learn the presented subject.

You are given an educational subject: {topic} 

Choose data and a chart that genuinely illustrate the concept, so a student can actually learn it from the figure.

After deciding the chart type and data, formulate the task containing these in order to be understood by a system that will generate the chart based on your decisions.

Reason briefly (2-3 sentences) and then write your answer in a JSON format.
"""

def generate_spec(topic):
    messages = [
        {"role": "user", "content": SPEC_PROMPT.format(topic=topic)}
    ]
    
    structured_outputs = StructuredOutputsParams(json=SpecificationsOutput.model_json_schema())

    params = SamplingParams(
            temperature = 0.7,
            top_p = 0.95,
            max_tokens = 4000,
            structured_outputs = structured_outputs
    )
    
    outputs = image_generation.model.chat(messages, params)
    
    try:
        return json.loads(outputs[0].outputs[0].text)
    except json.JSONDecodeError:
        return None