Update app.py

app.py

@@ -14,22 +14,11 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import transformers
 import torch
 
-model  = "tiiuae/falcon-40b-instruct"
 
-# model_name = 'google/flan-t5-base'
-# model = T5ForConditionalGeneration.from_pretrained(model_name, device_map='auto', offload_folder="offload")
-# tokenizer = AutoTokenizer.from_pretrained(model_name)
-# print('flan read')
+model_name = 'google/flan-t5-base'
+model = T5ForConditionalGeneration.from_pretrained(model_name, device_map='auto', offload_folder="offload")
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-tokenizer = AutoTokenizer.from_pretrained(model)
-pipeline = transformers.pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    torch_dtype=torch.bfloat16,
-    trust_remote_code=True,
-    device_map="auto",
-)
 
 
 ST_name = 'sentence-transformers/sentence-t5-base'
@@ -49,37 +38,18 @@ def get_context(query_text):
     return context
 
 def local_query(query, context):
-    # t5query = """Using the available context, please answer the question. 
-    # If you aren't sure please say i don't know.
-    # Context: {}
-    # Question: {}
-    # """.format(context, query)
-
-    # inputs = tokenizer(t5query, return_tensors="pt")
-
-    # outputs = model.generate(**inputs, max_new_tokens=20)
- 
-    # return tokenizer.batch_decode(outputs, skip_special_tokens=True)
-
-    context_query = """Using the available context, please answer the question. 
+    t5query = """Using the available context, please answer the question. 
     If you aren't sure please say i don't know.
     Context: {}
     Question: {}
     """.format(context, query)
-    
-    sequences = pipeline(
-    context_query,
-    max_length=200,
-    do_sample=True,
-    top_k=10,
-    num_return_sequences=1,
-    eos_token_id=tokenizer.eos_token_id,
-    )
-    
-    # for seq in sequences:
-    #     print(f"Result: {seq['generated_text']}")
 
-    return seq['generated_text']
+    inputs = tokenizer(t5query, return_tensors="pt")
+
+    outputs = model.generate(**inputs, max_new_tokens=20)
+ 
+    return tokenizer.batch_decode(outputs, skip_special_tokens=True)
+