ArunAIML's picture
bug issue resolved
1b2a00c
raw
history blame
771 Bytes
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from langchain_nvidia_ai_endpoints import ChatNVIDIA
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
import gradio as gr
import os
os.environ["NVIDIA_API_KEY"] = "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"
prompt = ChatPromptTemplate.from_messages([("system", "You are a helpful AI assistant named Fred."), ("user", "{input}")])
llm = ChatNVIDIA(model="mixtral_8x7b")
chain = prompt | llm | StrOutputParser()
def chat(prompt, history):
output = ""
for chunk in chain.stream({"input": prompt}):
output += chunk
yield output
demo = gr.ChatInterface(chat).queue()
demo.launch()