Vitrous commited on
Commit
ef40512
·
verified ·
1 Parent(s): 2233781

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +12 -3
app.py CHANGED
@@ -2,7 +2,7 @@ import gradio as gr
2
  import plotly.express as px
3
  import os
4
  import torch
5
- from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
6
 
7
  # Check if CUDA is available and set device accordingly
8
  device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -25,14 +25,23 @@ def hermes_model():
25
 
26
 
27
  def blender_model():
 
28
  tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot-400M-distill")
29
- model = AutoModelForCausalLM.from_pretrained("facebook/blenderbot-400M-distill", low_cpu_mem_usage=True, device_map="auto")
30
  return model, tokenizer
31
 
32
  model, tokenizer = blender_model()
33
 
34
- # Function to generate a response from the model
35
  def chat_response(msg_prompt: str) -> str:
 
 
 
 
 
 
 
 
 
 
36
  """
37
  Generates a response from the model given a prompt.
38
 
 
2
  import plotly.express as px
3
  import os
4
  import torch
5
+ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BlenderbotForConditionalGeneration
6
 
7
  # Check if CUDA is available and set device accordingly
8
  device = "cuda" if torch.cuda.is_available() else "cpu"
 
25
 
26
 
27
  def blender_model():
28
+ model = BlenderbotForConditionalGeneration.from_pretrained("facebook/blenderbot-400M-distill")
29
  tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot-400M-distill")
 
30
  return model, tokenizer
31
 
32
  model, tokenizer = blender_model()
33
 
 
34
  def chat_response(msg_prompt: str) -> str:
35
+ inputs = tokenizer([UTTERANCE], return_tensors="pt")
36
+ reply_ids = model.generate(**inputs)
37
+ outputs = tokenizer.batch_decode(reply_ids, skip_special_tokens=True)[0])
38
+
39
+ return outputs
40
+ except Exception as e:
41
+ return str(e)
42
+
43
+ # Function to generate a response from the model
44
+ def chat_responses(msg_prompt: str) -> str:
45
  """
46
  Generates a response from the model given a prompt.
47