Update app.py

app.py

@@ -1,52 +1,77 @@
 import gradio as gr
-from transformers import pipeline
+import requests
+import os
 
-# Initialize pipelines for text generation and speech recognition
-text_generation_pipeline = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2")
-speech_recognition_pipeline = pipeline("automatic-speech-recognition", model="openai/whisper-large-v3-turbo")
+# Set up the Hugging Face API key (ensure you've set this as an environment variable)
+api_key = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
-# Function to transcribe audio using the Whisper pipeline
+
+# 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-7B-Instruct-v0.2"
+
+# Set up headers for API requests
+headers = {"Authorization": f"Bearer {api_key}"}
+
+# Function to query the Hugging Face Whisper model for audio transcription
 def transcribe_audio(audio_file):
-    try:
-        # Use the Whisper pipeline for automatic speech recognition
-        transcription = speech_recognition_pipeline(audio_file)
-        return transcription.get("text", "Transcription not available.")
-    except Exception as e:
-        return f"Error in transcription: {e}"
-
-# Function to generate Mermaid.js code using the Mistral pipeline
+    with open(audio_file, "rb") as f:
+        data = f.read()
+    response = requests.post(WHISPER_API_URL, headers=headers, 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):
-    try:
-        # Use the text-generation pipeline to generate Mermaid.js code
-        result = text_generation_pipeline(
-            prompt, 
-            max_length=256, 
-            temperature=0.7, 
-            num_return_sequences=1
-        )
-        return result[0]["generated_text"].strip() if result else "No Mermaid.js code generated."
-    except Exception as e:
-        return f"Error in Mermaid.js generation: {e}"
-
-# Function to process input (text, audio, or both)
+    # 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}"
+
+# Function to process text, audio, or both inputs
 def process_input(input_type, text_input, audio_input):
-    try:
-        if input_type == "Audio" and audio_input:
-            transcription = transcribe_audio(audio_input)
-            return generate_mermaid_code(transcription)
+    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:
-            return generate_mermaid_code(text_input)
+    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:
-            transcription = transcribe_audio(audio_input)
-            combined_input = f"{text_input} and {transcription}"
-            return generate_mermaid_code(combined_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 "Please provide valid input."
-    except Exception as e:
-        return f"Error: {e}"
+    else:
+        return "No valid input provided."
 
 # Set up the Gradio interface
 iface = gr.Interface(
@@ -54,18 +79,14 @@ iface = gr.Interface(
     inputs=[
         gr.Radio(["Text", "Audio", "Text and Audio"], label="Input Type", value="Text"),
         gr.Textbox(lines=10, label="Text Input", placeholder="Enter task flow description here..."),
-        gr.Audio(type="filepath", label="Audio Input"),
+        gr.Audio(type="filepath", label="Audio Input")
     ],
     outputs=[
-        gr.Textbox(lines=20, label="Generated Mermaid.js Code"),
+        gr.Textbox(lines=20, label="Generated Mermaid.js Code")
     ],
     title="Mermaid.js Generator",
-    description=(
-        "Provide text, audio, or both."
-        "Mermaid.js code will be generated for text or audio input, or their combination."
-    ),
+    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
-if __name__ == "__main__":
-    iface.launch()
+iface.launch()