Update app.py

app.py

@@ -10,11 +10,26 @@ import uuid
 from sentence_transformers import SentenceTransformer
 import os
 
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import transformers
+import torch
+
+model_name  = "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)
-print('flan read')
+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'
@@ -34,17 +49,37 @@ def get_context(query_text):
     return context
 
 def local_query(query, context):
-    t5query = """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)
+
+    # 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. 
     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']}")
 
-    inputs = tokenizer(t5query, return_tensors="pt")
-
-    outputs = model.generate(**inputs, max_new_tokens=20)
- 
-    return tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    return seq['generated_text']