Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| import torch | |
| import re | |
| from gtts import gTTS | |
| import os | |
| import logging | |
| # Set up logging | |
| logging.basicConfig(filename='app.log', level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') | |
| # Function to set up the model and tokenizer | |
| def setup_model(model_name): | |
| logging.info('Setting up model and tokenizer.') | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_name, | |
| device_map="auto", | |
| trust_remote_code=False, | |
| revision="main" | |
| ) | |
| tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True) | |
| model.eval() | |
| logging.info('Model and tokenizer setup completed.') | |
| return model, tokenizer | |
| # Function to generate a response from the model | |
| def generate_response(model, tokenizer, prompt, max_new_tokens=140): | |
| logging.info('Generating response for the prompt.') | |
| inputs = tokenizer(prompt, return_tensors="pt") | |
| outputs = model.generate(input_ids=inputs["input_ids"].to("cuda"), max_new_tokens=max_new_tokens) | |
| response = tokenizer.batch_decode(outputs)[0] | |
| # Extract only the response part (assuming everything after the last newline belongs to the response) | |
| response_parts = response.split("\n") | |
| logging.info('Response generated.') | |
| return response_parts[-1] # Return the last element (response) | |
| # Function to remove various tags using regular expressions | |
| def remove_tags(text): | |
| logging.info('Removing tags from the text.') | |
| # Combine multiple tag removal patterns for broader coverage | |
| tag_regex = r"<[^>]*>" # Standard HTML tags | |
| custom_tag_regex = r"<.*?>|\[.*?\]|{\s*?\(.*?\)\s*}" # Custom, non-standard tags (may need adjustments) | |
| all_tags_regex = f"{tag_regex}|{custom_tag_regex}" # Combine patterns | |
| cleaned_text = re.sub(all_tags_regex, "", text) | |
| logging.info('Tags removed.') | |
| return cleaned_text | |
| # Function to generate the audio file | |
| def text_to_speech(text, filename="response.mp3"): | |
| logging.info('Generating speech audio file.') | |
| tts = gTTS(text) | |
| tts.save(filename) | |
| logging.info('Speech audio file saved.') | |
| return filename | |
| # Main function for the Gradio app | |
| def main(comment): | |
| logging.info('Main function triggered.') | |
| instructions_string = ( | |
| "virtual marketer assistant, communicates in business, focused on services, " | |
| "escalating to technical depth upon request. It reacts to feedback aptly and ends responses " | |
| "with its signature Mr.jon will tailor the length of its responses to match the individual's comment, " | |
| "providing concise acknowledgments to brief expressions of gratitude or feedback, thus keeping the interaction natural and supportive.\n" | |
| ) | |
| model_name = "TheBloke/Mistral-7B-Instruct-v0.2-GPTQ" | |
| try: | |
| model, tokenizer = setup_model(model_name) | |
| if comment: | |
| prompt_template = lambda comment: f"[INST] {instructions_string} \n{comment} \n[/INST]" | |
| prompt = prompt_template(comment) | |
| response = generate_response(model, tokenizer, prompt) | |
| # Apply tag removal before displaying the response | |
| response_without_tags = remove_tags(response) | |
| # Remove the "[/INST]" string at the end (assuming it's always present) | |
| response_without_inst = response_without_tags.rstrip("[/INST]") | |
| # Generate and return the response and the audio file | |
| audio_file = text_to_speech(response_without_inst) | |
| logging.info('Response and audio file generated.') | |
| return response_without_inst, audio_file | |
| else: | |
| logging.warning('No comment entered.') | |
| return "Please enter a comment to generate a response.", None | |
| except Exception as e: | |
| logging.error(f'Error occurred: {str(e)}') | |
| return "An error occurred. Please try again later.", None | |
| iface = gr.Interface( | |
| fn=main, | |
| inputs=gr.Textbox(lines=2, placeholder="Enter a comment..."), | |
| outputs=["text", "file"], | |
| title="Virtual Marketer Assistant", | |
| description="Enter a comment and get a response from the virtual marketer assistant. Download the response as an MP3 file." | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch(share=True) | |