Update app.py

app.py

@@ -1,26 +1,25 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoProcessor, AutoModel, BitsAndBytesConfig
+from transformers import AutoProcessor, AutoModel, BitsAndBytesConfig
 import scipy.io.wavfile as wav
 import numpy as np
+import os
+import openai
+
+# Set OpenAI API key (preferably use environment variables)
+os.environ["OPENAI_API_KEY"] = "sk-proj-hx8b2Z3ja8dyzDH1DtCOT3BlbkFJPV804EO36VB4LbflIUa7"
+openai.api_key = os.getenv("OPENAI_API_KEY")
 
 # Check for GPU availability
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Quantization config
+# Quantization config for speech model
 if device == "cuda":
     quantization_config = BitsAndBytesConfig(load_in_8bit=True)
 else:
     quantization_config = None
 
-# Load models and tokenizers
-story_tokenizer = AutoTokenizer.from_pretrained("starmpcc/Asclepius-Llama3-8B", use_fast=False)
-story_model = AutoModelForCausalLM.from_pretrained(
-    "starmpcc/Asclepius-Llama3-8B", 
-    device_map="auto" if device == "cuda" else None,
-    quantization_config=quantization_config
-)
-
+# Load speech model and processor
 speech_processor = AutoProcessor.from_pretrained("suno/bark-small")
 speech_model = AutoModel.from_pretrained(
     "suno/bark-small", 
@@ -28,27 +27,34 @@ speech_model = AutoModel.from_pretrained(
     quantization_config=quantization_config
 )
 
-# Move models to the appropriate device
-story_model.to(device)
+# Move speech model to the appropriate device
 speech_model.to(device)
 
-# Function to generate story
+# Function to generate story using GPT
 def generate_story(prompt):
-    model_input = f"""You are an intelligent clinical language model.
-    Below is a snippet of patient's discharge summary and a following instruction from healthcare professional.
-    Write a response that appropriately completes the instruction.
-    The response should provide the accurate answer to the instruction, while being concise.
-    [Discharge Summary Begin]
+    model_input = f"""You are a creative and educational storyteller for school-going children.
+    Your task is to create an engaging, age-appropriate story that both entertains and teaches valuable lessons.
+    Use the following prompt as inspiration for your story, but feel free to be imaginative and expand upon it.
+    Remember to include educational elements that children can learn from, such as historical facts, scientific concepts, moral lessons, or cultural insights.
+    
+    [Story Prompt Begin]
     {prompt}
-    [Discharge Summary End]
+    [Story Prompt End]
+    
     [Instruction Begin]
-    Generate a short story based on this prompt.
-    [Instruction End] 
+    Generate a short, creative, and educational story based on this prompt. The story should be suitable for school-going children, entertaining, and contain clear learning points.
+    [Instruction End]
     """
     
-    input_ids = story_tokenizer(model_input, return_tensors="pt").input_ids.to(device)
-    output = story_model.generate(input_ids, max_length=500)
-    return story_tokenizer.decode(output[0], skip_special_tokens=True)
+    response = openai.ChatCompletion.create(
+        model="gpt-3.5-turbo",  # or "gpt-4" if you have access
+        messages=[
+            {"role": "system", "content": "You are a skilled storyteller who creates educational and engaging stories for children."},
+            {"role": "user", "content": model_input}
+        ]
+    )
+    
+    return response.choices[0].message['content']
 
 # Function to generate speech
 def generate_speech(text, speaker="v2/en_speaker_6"):
@@ -75,7 +81,7 @@ iface = gr.Interface(
         gr.Audio(label="Generated Speech")
     ],
     title="Story Generator and Text-to-Speech",
-    description="Enter a prompt to generate a story, then convert it to speech."
+    description="Enter a prompt to generate a story using GPT, then convert it to speech."
 )
 
 # Launch the app