speech_to_text.py

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import gradio as gr

model_name = "mistralai/Mistral-7B-Instruct-v0.1"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")

chat_history = [{"role": "system", "content": "You are a helpful assistant."}]

def chat(user_input):
    # 사용자 입력을 채팅 기록에 추가
    chat_history.append({"role": "user", "content": user_input})
    
    # 모델에 채팅 기록 전달
    inputs = tokenizer([message['content'] for message in chat_history], return_tensors="pt", padding=True).to("cuda" if torch.cuda.is_available() else "cpu")
    
    # 모델로 응답 생성
    outputs = model.generate(**inputs, max_length=200)
    
    # 생성된 응답
    bot_reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # 봇 응답을 채팅 기록에 추가
    chat_history.append({"role": "assistant", "content": bot_reply})
    
    return bot_reply

demo = gr.ChatInterface(fn=chat, type='messages', title='이우진의 챗봇')
demo.launch(share=True)