Update app.py

app.py

@@ -1,81 +1,51 @@
 import gradio as gr
 import requests
 import os
+from huggingface_hub import InferenceClient
 
 # Set up the Hugging Face API key (ensure you've set this as an environment variable)
 api_key = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
-
 # API URLs
 WHISPER_API_URL = "https://api-inference.huggingface.co/models/openai/whisper-large-v3-turbo"
-#MISTRAL_API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-Nemo-Instruct-2407"
-MISTRAL_API_URL = "https://api-inference.huggingface.co/models/deepseek-ai/DeepSeek-R1"
 
-deepseek-ai/DeepSeek-R1
-# Set up headers for API requests
-headers = {"Authorization": f"Bearer {api_key}"}
+# Set up inference client for DeepSeek-R1
+client = InferenceClient(
+    provider="together",
+    api_key=api_key
+)
 
-# Function to query the Hugging Face Whisper model for audio transcription
 def transcribe_audio(audio_file):
     with open(audio_file, "rb") as f:
         data = f.read()
-    response = requests.post(WHISPER_API_URL, headers=headers, data=data)
+    response = requests.post(WHISPER_API_URL, headers={"Authorization": f"Bearer {api_key}"}, data=data)
     if response.status_code == 200:
         return response.json().get("text", "Transcription not available.")
     else:
         return f"Error: {response.status_code}, {response.text}"
 
-# Function to query the Mistral model to generate Mermaid.js code
 def generate_mermaid_code(prompt):
-    # mermaid_prompt = f"Use the appropriate diagram type (Use Case Diagram, Flowchart, Sequence Diagram, Entity-Relationship (ER) Diagram,State Diagram, Pie Chart etc.) based on the context.\n" \
-    #                  "Generate a valid, syntactically correct MermaidJS diagram code for the following: {prompt}"
-    mermaid_prompt = f"Generate a valid MermaidJS diagram code for the following: {prompt}"
-
-
-    # Prepare the payload (input for the model)
-    payload = {
-        "inputs": mermaid_prompt,
-        "parameters": {
-            "max_length": 256,
-            "temperature": 0.7
-        }
-    }
-
-    # Send the request to the Mistral API
-    response = requests.post(MISTRAL_API_URL, headers=headers, json=payload)
-
-    # Check if the request was successful
-    if response.status_code == 200:
-        result = response.json()
-        # Extract the generated Mermaid.js code
-        return result[0]['generated_text'].strip()
-    else:
-        return f"Error: {response.status_code}, {response.text}"
+    messages = [{"role": "user", "content": f"Generate a valid MermaidJS diagram code for the following: {prompt}"}]
+    completion = client.chat.completions.create(
+        model="deepseek-ai/DeepSeek-R1",
+        messages=messages,
+        max_tokens=500
+    )
+    return completion.choices[0].message['content'].strip()
 
-# Function to process text, audio, or both inputs
 def process_input(input_type, text_input, audio_input):
     if input_type == "Audio" and audio_input is not None:
-        # Transcribe audio
         transcription = transcribe_audio(audio_input)
-        # Generate Mermaid.js code
         return generate_mermaid_code(transcription)
-
     elif input_type == "Text" and text_input:
-        # Generate Mermaid.js code directly from text input
         return generate_mermaid_code(text_input)
-
     elif input_type == "Text and Audio" and text_input and audio_input is not None:
-        # Transcribe audio
         transcription = transcribe_audio(audio_input)
-        # Combine text input and transcription
         combined_input = f"{text_input} and {transcription}"
-        # Generate Mermaid.js code
         return generate_mermaid_code(combined_input)
-
     else:
         return "No valid input provided."
 
-# Set up the Gradio interface
 iface = gr.Interface(
     fn=process_input,
     inputs=[
@@ -90,5 +60,4 @@ iface = gr.Interface(
     description="Provide text, audio, or both. Mermaid.js code will be generated based on the text or audio input, or their combination."
 )
 
-# Launch the Gradio app
 iface.launch()