deanoreese commited on
Commit
ba3a92a
·
verified ·
1 Parent(s): e0bc044

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -6
app.py CHANGED
@@ -15,17 +15,18 @@ import os
15
 
16
  class BasicAgent:
17
  def __init__(self):
18
- print(" Using Falcon-RW-1B for fast runs...")
19
  self.generator = pipeline(
20
- "text-generation",
21
- model="tiiuae/falcon-rw-1b",
22
  token=os.getenv("HF_API_KEY")
23
  )
24
 
25
  def __call__(self, question: str) -> str:
26
- prompt = f"Answer this question accurately: {question}"
27
- result = self.generator(prompt, max_length=100, do_sample=False)
28
- return result[0]['generated_text'].replace(prompt, '').strip()
 
29
 
30
  def run_and_submit_all( profile: gr.OAuthProfile | None):
31
  """
 
15
 
16
  class BasicAgent:
17
  def __init__(self):
18
+ print(" Using flan-t5-base for fast generation...")
19
  self.generator = pipeline(
20
+ "text2text-generation", # Note: different pipeline type for T5-style models
21
+ model="google/flan-t5-base",
22
  token=os.getenv("HF_API_KEY")
23
  )
24
 
25
  def __call__(self, question: str) -> str:
26
+ prompt = f"Answer this question: {question}"
27
+ result = self.generator(prompt, max_new_tokens=50, truncation=True)
28
+ return result[0]['generated_text'].strip()
29
+
30
 
31
  def run_and_submit_all( profile: gr.OAuthProfile | None):
32
  """