nvdia llm mistral

app.py

@@ -1,21 +1,24 @@
 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
 
-model_id = "gpt2"
+os.environ["NVIDIA_API_KEY"] = "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"
 
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id)
-model.to_bettertransformer()
+prompt = ChatPromptTemplate.from_messages([("system", "You are a helpful AI assistant named Fred."), ("user", "{input}")])
 
-pipe = pipeline("text-generation", model=model, tokenizer= tokenizer)
+llm = ChatNVIDIA(model="mixtral_8x7b")
+chain = prompt | llm | StrOutputParser()
 
-def gpt(prompt, top_k, penalty_alpha):
-    return pipe(prompt, top_k=top_k, penalty_alpha=penalty_alpha)[0]["generated_text"]
 
-gr.Interface(
-    gpt,
-    ["text",gr.Slider(minimum=0, maximum=50, step=1,label="Top_k"),gr.Slider(minimum=0.1, maximum=1.0,label="penalty_alpha")],
-    "text",
-    title= "Arun's GPT chatbot",
-    description = "This is Arun's experimental GPT interface exposing gpt2, feel free to experiment"
-).launch()
+def chat(prompt, history):
+
+    for chunk in chain.stream({"input": prompt}):
+        yield chunk.content
+
+gr.Chat
+demo = gr.ChatInterface(chat).queue()
+
+demo.launch()

requirements.txt

@@ -2,3 +2,5 @@ transformers
 gradio
 torch
 optimum
+langchain
+langchain-nvidia-ai-endpoints