electblake commited on
Commit
fab4800
·
verified ·
1 Parent(s): ca5bf81

Constrain XLSX inference to structured table output

Browse files

Use llama.cpp JSON-schema constrained generation for columns and rows, then construct the Gradio dataframe directly. This fixes malformed free-form CSV output on ZeroGPU.

Files changed (1) hide show
  1. app.py +27 -5
app.py CHANGED
@@ -1,4 +1,4 @@
1
- from io import StringIO
2
  from pathlib import Path
3
  import tomllib
4
 
@@ -190,8 +190,9 @@ def generate_xlsx_table(
190
  {
191
  "role": "system",
192
  "content": (
193
- f"{system_prompt}\n\nReturn only valid CSV with one header row. "
194
- "Do not wrap the CSV in a code fence or add prose."
 
195
  ),
196
  },
197
  {
@@ -205,13 +206,34 @@ def generate_xlsx_table(
205
  ]
206
  completion = model.create_chat_completion(
207
  messages=messages,
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
208
  max_tokens=1024,
209
  temperature=0.6,
210
  top_p=0.95,
211
  top_k=20,
212
  )
213
- response = completion["choices"][0]["message"]["content"].rsplit("</think>", 1)[-1].strip()
214
- return pd.read_csv(StringIO(response))
215
 
216
 
217
  CSS = """
 
1
+ import json
2
  from pathlib import Path
3
  import tomllib
4
 
 
190
  {
191
  "role": "system",
192
  "content": (
193
+ f"{system_prompt}\n\nReturn only a JSON table object with a columns "
194
+ "array and a rows array. Every row must contain one string value per "
195
+ "column."
196
  ),
197
  },
198
  {
 
206
  ]
207
  completion = model.create_chat_completion(
208
  messages=messages,
209
+ response_format={
210
+ "type": "json_object",
211
+ "schema": {
212
+ "type": "object",
213
+ "properties": {
214
+ "columns": {
215
+ "type": "array",
216
+ "items": {"type": "string"},
217
+ },
218
+ "rows": {
219
+ "type": "array",
220
+ "items": {
221
+ "type": "array",
222
+ "items": {"type": "string"},
223
+ },
224
+ },
225
+ },
226
+ "required": ["columns", "rows"],
227
+ "additionalProperties": False,
228
+ },
229
+ },
230
  max_tokens=1024,
231
  temperature=0.6,
232
  top_p=0.95,
233
  top_k=20,
234
  )
235
+ response = json.loads(completion["choices"][0]["message"]["content"])
236
+ return pd.DataFrame(response["rows"], columns=response["columns"])
237
 
238
 
239
  CSS = """