Update app.py

app.py

@@ -2,17 +2,22 @@ import gradio as gr
 import json
 import spaces
 import torch
-from unsloth import FastLanguageModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
 
 # ── Load model once at startup ──────────────────────────────
-MODEL_ID = "suneeldk/json-extract" 
+BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
+LORA_MODEL = "YOUR_USERNAME/json-extract"  # ← change this
 
-model, tokenizer = FastLanguageModel.from_pretrained(
-    MODEL_ID,
-    max_seq_length=2048,
-    load_in_4bit=True,
+tokenizer = AutoTokenizer.from_pretrained(LORA_MODEL)
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.float16,
+    device_map="auto",
 )
-FastLanguageModel.for_inference(model)
+model = PeftModel.from_pretrained(base_model, LORA_MODEL)
+model.eval()
 
 
 # ── Inference function ──────────────────────────────────────
@@ -29,7 +34,7 @@ def extract(text, schema_text):
         return "Invalid JSON schema. Please check the format."
 
     prompt = f"### Input: {text}\n### Schema: {json.dumps(schema)}\n### Output:"
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
     with torch.no_grad():
         outputs = model.generate(
@@ -38,7 +43,6 @@ def extract(text, schema_text):
             temperature=0.1,
             do_sample=True,
             pad_token_id=tokenizer.eos_token_id,
-            use_cache=False,
         )
 
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)