Update generator.py to electra small

generator.py

@@ -37,11 +37,12 @@ import streamlit as st
 def load_model():
   hfm = pickle.load(open('hfmodel.sav','rb'))
   hft = T5TokenizerFast.from_pretrained("t5-base")
-  tok = att.from_pretrained("ahotrod/albert_xxlargev1_squad2_512")
-  model = pickle.load(open('model.sav','rb'))
-  return hfm, hft,tok, model
+  # tok = att.from_pretrained("")
+  model = pickle.load(open('electra_model.sav','rb'))
+  # return hfm, hft,tok, model
+  return hfm, hft, model
 
-hfmodel, hftokenizer, tokenizer, model = load_model()
+hfmodel, hftokenizer, model = load_model()
 
 def run_model(input_string, **generator_args):
   generator_args = {
@@ -67,29 +68,29 @@ def run_model(input_string, **generator_args):
 # al_tokenizer = pickle.load(open('models/al_tokenizer.sav', 'rb'))
 def QA(question, context):
   # model_name="deepset/electra-base-squad2"
-  # nlp = pipeline("question-answering",model=al_model,tokenizer=al_tokenizer)
-  # format = {
-  #     'question':question,
-  #     'context':context
-  # }
-  # res = nlp(format)
-  # output = f"{question}\n{string.capwords(res['answer'])}\tscore : [{res['score']}] \n"
-  # return output
-  inputs = tokenizer(question, context, return_tensors="pt")
-  # Run the model, the deepset way
-  with torch.no_grad():
-    output = model(**inputs)
-  start_score = output.start_logits
-  end_score = output.end_logits
-  #Get the rel scores for the context, and calculate the most probable begginign using torch
-  start = torch.argmax(start_score)
-  end = torch.argmax(end_score)
-  #cinvert tokens to strings
-  # output = tokenizer.decode(input_ids[start:end+1], skip_special_tokens=True)
-  predict_answer_tokens = inputs.input_ids[0, start : end + 1]
-  output = tokenizer.decode(predict_answer_tokens, skip_special_tokens=True)
-  output = string.capwords(output)
-  return f"Q. {question} \n Ans. {output}"
+  nlp = pipeline("question-answering",model=model)
+  format = {
+      'question':question,
+      'context':context
+  }
+  res = nlp(format)
+  output = f"{question}\n{string.capwords(res['answer'])}\tscore : [{res['score']}] \n"
+  return output
+  # inputs = tokenizer(question, context, return_tensors="pt")
+  # # Run the model, the deepset way
+  # with torch.no_grad():
+  #   output = model(**inputs)
+  # start_score = output.start_logits
+  # end_score = output.end_logits
+  # #Get the rel scores for the context, and calculate the most probable begginign using torch
+  # start = torch.argmax(start_score)
+  # end = torch.argmax(end_score)
+  # #cinvert tokens to strings
+  # # output = tokenizer.decode(input_ids[start:end+1], skip_special_tokens=True)
+  # predict_answer_tokens = inputs.input_ids[0, start : end + 1]
+  # output = tokenizer.decode(predict_answer_tokens, skip_special_tokens=True)
+  # output = string.capwords(output)
+  # return f"Q. {question} \n Ans. {output}"
 # QA("What was the first C program","The first prgram written in C was Hello World")
 
 def gen_question(inputs):