Update app.py

app.py

@@ -15,17 +15,19 @@ model.eval()
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-def generate_text(prompt):
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-    tokens = model.generate(input_ids, max_length=22)
-    return tokenizer.decode(tokens[0].tolist(), skip_special_tokens=False)
+def generate_text(question):
+    prompt = f'Q: {question}\nA: '
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    tokens = model.generate(input_ids, max_length=50, pad_token_id=tokenizer.eos_token_id)
+    response = tokenizer.decode(tokens[0], skip_special_tokens=True)
+    return response.split('\nA: ')[-1] 
 
 iface = gr.Interface(
+    gr.Markdown(title),
+    gr.Markdown(description),
     fn=generate_text,
     inputs=gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
-    outputs="text",
-    title=title,
-    description=description
+    outputs="text"
 )
 
 iface.launch()