Update app.py

app.py

@@ -1,19 +1,18 @@
 import gradio as gr
 import requests
 import os
-from transformers import pipeline
+import json
 
 # Set up the Hugging Face API key (ensure you've set this as an environment variable)
 api_key = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
-# API URL for Whisper (audio transcription)
+# API URLs
 WHISPER_API_URL = "https://api-inference.huggingface.co/models/openai/whisper-large-v3-turbo"
+# MISTRAL_API_URL (DeepSeek API call for generating Mermaid code)
+MISTRAL_API_URL = "https://huggingface.co/api/inference-proxy/together/v1/chat/completions"
 
-# Set up headers for the Whisper API request
-headers = {"Authorization": f"Bearer {api_key}"}
-
-# Load the DeepSeek-R1-Distill-Qwen-1.5B model using Hugging Face's pipeline
-pipe = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B")
+# Set up headers for API requests
+headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
 
 # Function to query the Hugging Face Whisper model for audio transcription
 def transcribe_audio(audio_file):
@@ -25,33 +24,47 @@ def transcribe_audio(audio_file):
     else:
         return f"Error: {response.status_code}, {response.text}"
 
-# Function to generate Mermaid.js code using the DeepSeek model (DeepSeek-R1-Distill-Qwen-1.5B)
+# Function to query the Hugging Face API to generate Mermaid.js code
 def generate_mermaid_code(prompt):
-    # Instruction included in the prompt to guide DeepSeek to generate valid MermaidJS code
-    deepseek_prompt = f"Generate a valid MermaidJS diagram code for the following: {prompt}"
-    
-    # Using the DeepSeek model pipeline for text generation
-    response = pipe([{"role": "user", "content": deepseek_prompt}])
-    return response[0]["generated_text"].strip()
+    # Define the payload to send to the Hugging Face API
+    mermaid_prompt = f"Generate a valid MermaidJS diagram code for the following: {prompt}"
+
+    payload = {
+        "model": "deepseek-ai/DeepSeek-R1",
+        "messages": [{"role": "user", "content": mermaid_prompt}],
+        "max_tokens": 500,
+        "stream": False
+    }
+
+    # Send the request to the Hugging Face API
+    response = requests.post(MISTRAL_API_URL, headers=headers, data=json.dumps(payload))
+
+    # Check if the request was successful
+    if response.status_code == 200:
+        result = response.json()
+        # Extract the generated Mermaid.js code from the response
+        return result['choices'][0]['message']['content'].strip()
+    else:
+        return f"Error: {response.status_code}, {response.text}"
 
 # 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 using the Whisper API
+        # Transcribe audio
         transcription = transcribe_audio(audio_input)
-        # Generate Mermaid.js code from transcription using DeepSeek-R1-Distill-Qwen-1.5B
+        # Generate Mermaid.js code from the transcription
         return generate_mermaid_code(transcription)
 
     elif input_type == "Text" and text_input:
-        # Generate Mermaid.js code directly from text input using DeepSeek-R1-Distill-Qwen-1.5B
+        # 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 using the Whisper API
+        # Transcribe audio
         transcription = transcribe_audio(audio_input)
         # Combine text input and transcription
         combined_input = f"{text_input} and {transcription}"
-        # Generate Mermaid.js code using DeepSeek-R1-Distill-Qwen-1.5B
+        # Generate Mermaid.js code from the combined input
         return generate_mermaid_code(combined_input)
 
     else: