Update app.py

app.py

@@ -99,29 +99,27 @@ tokenizer.save_pretrained("/fine_tuned_model")
 
 
 import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
+# Initialize fine-tuned model and tokenizer
 model_dir= "/fine_tuned_model"
-tokenizer = GPT2Tokenizer.from_pretrained(model_dir)
-model = GPT2LMHeadModel.from_pretrained(model_dir)
-
-# Load chromadb collection to pass as context
-documents = collection.get(["documents"])
-
-# Define function to answer questions using the fine-tuned model and ChromaDB collection
-def answer_question(question):
-    # Concatenate document contents to create context
-    context = " ".join(doc["content"] for doc in documents)
-
-    # Append question to the context
-    input_text = f"Question: {question} Context: {context} Answer:"
-
-    # Generate answer using the model
-    input_ids = tokenizer.encode(input_text, return_tensors="pt")
-    generated = model.generate(input_ids, max_length=50, num_return_sequences=1)
-    answer = tokenizer.decode(generated[0], skip_special_tokens=True)
-
-
-gr.Interface(fn=answer_question, inputs="text", outputs="text").launch()
-
+tokenizer = AutoTokenizer.from_pretrained(model_dir)
+model = AutoModelForCausalLM.from_pretrained(model_dir)
+
+# Create a text generation pipeline
+text_gen = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+# Define question_answering function
+def question_answer(question):
+  generated = text_gen(question, max_length=200, num_return_sequences=1)
+  generated_tokens = generated[0]['generated_text'].replace(question, "")
+  generated_token_ids = [int(token) for token in generated_tokens.strip().split()]
+  answer = tokenizer.decode(generated_token_ids)
+  return answer
+
+# Set up gradio interface
+iface = gr.Interface(fn = question_answer, inputs='text', outputs='text', title='Fine-tuned Enron Question Answering',
+                     description='Ask a question regarding the Enron case')
+iface.launch()
 """## Deploy the Gradio Interface in a Huggingface Space"""