aaliyaan commited on
Commit
8431cf4
·
1 Parent(s): f8a2755

Fixed model loading for Bloom

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -1,11 +1,11 @@
1
  import gradio as gr
2
- from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
3
  from PyPDF2 import PdfReader
4
 
5
  # Models and tokenizers setup
6
  models = {
7
  "Text Generator (Bloom)": {
8
- "model": AutoModelForSeq2SeqLM.from_pretrained("bigscience/bloom-560m"),
9
  "tokenizer": AutoTokenizer.from_pretrained("bigscience/bloom-560m"),
10
  },
11
  "PDF Summarizer (T5)": {
@@ -23,7 +23,7 @@ def generate_text(model_choice, input_text, max_tokens, temperature, top_p):
23
  model_info = models[model_choice]
24
  tokenizer = model_info["tokenizer"]
25
  model = model_info["model"]
26
-
27
  inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
28
  outputs = model.generate(
29
  **inputs, max_length=max_tokens, num_beams=5, early_stopping=True, temperature=temperature, top_p=top_p
 
1
  import gradio as gr
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
  from PyPDF2 import PdfReader
4
 
5
  # Models and tokenizers setup
6
  models = {
7
  "Text Generator (Bloom)": {
8
+ "model": AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m"),
9
  "tokenizer": AutoTokenizer.from_pretrained("bigscience/bloom-560m"),
10
  },
11
  "PDF Summarizer (T5)": {
 
23
  model_info = models[model_choice]
24
  tokenizer = model_info["tokenizer"]
25
  model = model_info["model"]
26
+
27
  inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
28
  outputs = model.generate(
29
  **inputs, max_length=max_tokens, num_beams=5, early_stopping=True, temperature=temperature, top_p=top_p