Update app.py

app.py

@@ -8,18 +8,17 @@ import json
 from huggingface_hub import InferenceClient
 import spaces
 
-# Initialize the InferenceClient
-client = InferenceClient(model="gpt-3.5-turbo")
+def ai_fix_json(model_id, json_data):
+    # Initialize the InferenceClient with the chosen model
+    client = InferenceClient(model=model_id)
 
-@spaces.GPU(duration=120)
-def ai_fix_json(json_data):
     prompt = f"Fix the following JSON data and make it valid:\n\n{json_data}\n\nFixed JSON:"
     
     try:
         if torch.cuda.is_available():
             response = client.text_generation(prompt, max_new_tokens=1024)
             fixed_json = response[0]['generated_text'].split("Fixed JSON:")[-1].strip()
-            return fixed_json, "JSON fixed using AI on GPU."
+            return fixed_json, "JSON fixed using AI on GPU with model: " + model_id
         else:
             raise RuntimeError("GPU not available, falling back to CPU.")
     
@@ -28,11 +27,11 @@ def ai_fix_json(json_data):
         print(f"Falling back to CPU due to: {gpu_error}")
         response = client.text_generation(prompt, max_new_tokens=1024)
         fixed_json = response[0]['generated_text'].split("Fixed JSON:")[-1].strip()
-        return fixed_json, "JSON fixed using AI on CPU."
+        return fixed_json, "JSON fixed using AI on CPU with model: " + model_id
 
-def process_file(uploaded_file):
+def process_file(model_id, uploaded_file):
     json_data = uploaded_file  # This is already the content of the file as a string
-    cleaned_json, message = ai_fix_json(json_data)
+    cleaned_json, message = ai_fix_json(model_id, json_data)
     
     try:
         parsed_data = json.loads(cleaned_json)
@@ -43,10 +42,11 @@ def process_file(uploaded_file):
 
 iface = gr.Interface(
     fn=process_file,
-    inputs=gr.File(label="Upload your JSON file"),
+    inputs=[gr.Textbox(label="Model ID (e.g., gpt-3.5-turbo)", default="gpt-3.5-turbo"),
+            gr.File(label="Upload your JSON file")],
     outputs=[gr.JSON(label="Fixed JSON"), "text", gr.File(label="Download cleaned JSON file")],
-    title="AI-Powered JSON Cleaner with GPU Support",
-    description="Upload a JSON file to automatically fix, remove duplicates, and download the cleaned version using AI with GPU/CPU fallback."
+    title="AI-Powered JSON Cleaner with Model Selection",
+    description="Upload a JSON file to automatically fix, remove duplicates, and download the cleaned version using AI with GPU/CPU fallback. Select any model from the Hugging Face Hub."
 )
 
 if __name__ == "__main__":